[model-gateway] convert workflow system to type-safe workflow data (#16970)

sgl-model-gateway/src/app_context.rs

@@ -8,7 +8,10 @@ use tracing::{debug, info};
 
 use crate::{
     config::RouterConfig,
-    core::{JobQueue, LoadMonitor, WorkerRegistry, WorkerService, UNKNOWN_MODEL_ID},
+    core::{
+        steps::workflow_data::AnyWorkflowData, JobQueue, LoadMonitor, WorkerRegistry,
+        WorkerService, UNKNOWN_MODEL_ID,
+    },
     data_connector::{
         create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage,
     },
@@ -26,9 +29,12 @@ use crate::{
     },
     tool_parser::ParserFactory as ToolParserFactory,
     wasm::{config::WasmRuntimeConfig, module_manager::WasmModuleManager},
-    workflow::WorkflowEngine,
+    workflow::{InMemoryStore, WorkflowEngine},
 };
 
+/// Type alias for the concrete workflow engine used in the application
+pub type AppWorkflowEngine = WorkflowEngine<AnyWorkflowData, InMemoryStore<AnyWorkflowData>>;
+
 /// Error type for AppContext builder
 #[derive(Debug)]
 pub struct AppContextBuildError(&'static str);
@@ -59,13 +65,21 @@ pub struct AppContext {
     pub configured_reasoning_parser: Option<String>,
     pub configured_tool_parser: Option<String>,
     pub worker_job_queue: Arc<OnceLock<Arc<JobQueue>>>,
-    pub workflow_engine: Arc<OnceLock<Arc<WorkflowEngine>>>,
+    pub workflow_engine: Arc<OnceLock<Arc<AppWorkflowEngine>>>,
     pub mcp_manager: Arc<OnceLock<Arc<McpManager>>>,
     pub wasm_manager: Option<Arc<WasmModuleManager>>,
     pub worker_service: Arc<WorkerService>,
     pub inflight_tracker: Arc<InFlightRequestTracker>,
 }
 
+impl std::fmt::Debug for AppContext {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("AppContext")
+            .field("router_config", &self.router_config)
+            .finish_non_exhaustive()
+    }
+}
+
 pub struct AppContextBuilder {
     client: Option<Client>,
     router_config: Option<RouterConfig>,
@@ -81,7 +95,7 @@ pub struct AppContextBuilder {
     conversation_item_storage: Option<Arc<dyn ConversationItemStorage>>,
     load_monitor: Option<Arc<LoadMonitor>>,
     worker_job_queue: Option<Arc<OnceLock<Arc<JobQueue>>>>,
-    workflow_engine: Option<Arc<OnceLock<Arc<WorkflowEngine>>>>,
+    workflow_engine: Option<Arc<OnceLock<Arc<AppWorkflowEngine>>>>,
     mcp_manager: Option<Arc<OnceLock<Arc<McpManager>>>>,
     wasm_manager: Option<Arc<WasmModuleManager>>,
 }
@@ -206,7 +220,10 @@ impl AppContextBuilder {
         self
     }
 
-    pub fn workflow_engine(mut self, workflow_engine: Arc<OnceLock<Arc<WorkflowEngine>>>) -> Self {
+    pub fn workflow_engine(
+        mut self,
+        workflow_engine: Arc<OnceLock<Arc<AppWorkflowEngine>>>,
+    ) -> Self {
         self.workflow_engine = Some(workflow_engine);
         self
     }

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

@@ -14,15 +14,19 @@ use tokio::sync::{mpsc, Semaphore};
 use tracing::{debug, error, info, warn};
 
 use crate::{
-    app_context::AppContext,
+    app_context::{AppContext, AppWorkflowEngine},
     config::{RouterConfig, RoutingMode},
     core::steps::{
+        create_external_worker_workflow_data, create_local_worker_workflow_data,
+        create_mcp_workflow_data, create_tokenizer_workflow_data,
+        create_wasm_registration_workflow_data, create_wasm_removal_workflow_data,
+        create_worker_removal_workflow_data, create_worker_update_workflow_data,
         McpServerConfigRequest, TokenizerConfigRequest, TokenizerRemovalRequest,
-        WasmModuleConfigRequest, WasmModuleRemovalRequest, WorkerRemovalRequest,
+        WasmModuleConfigRequest, WasmModuleRemovalRequest,
     },
     mcp::McpConfig,
     protocols::worker_spec::{JobStatus, WorkerConfigRequest, WorkerUpdateRequest},
-    workflow::{WorkflowContext, WorkflowEngine, WorkflowId, WorkflowInstanceId, WorkflowStatus},
+    workflow::{WorkflowId, WorkflowInstanceId, WorkflowStatus},
 };
 
 /// Job types for control plane operations
@@ -404,17 +408,11 @@ impl JobQueue {
                     .get()
                     .ok_or_else(|| "Workflow engine not initialized".to_string())?;
 
-                let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
+                let workflow_data =
-                // Convert Box to Arc for context storage
+                    create_wasm_registration_workflow_data(*config.clone(), Arc::clone(context));
-                let config_arc: Arc<WasmModuleConfigRequest> = Arc::new(*config.clone());
-                workflow_context.set_arc("wasm_module_config", config_arc);
-                workflow_context.set_arc("app_context", Arc::clone(context));
 
                 let instance_id = engine
-                    .start_workflow(
+                    .start_workflow(WorkflowId::new("wasm_module_registration"), workflow_data)
-                        WorkflowId::new("wasm_module_registration"),
-                        workflow_context,
-                    )
                     .await
                     .map_err(|e| {
                         format!("Failed to start WASM module registration workflow: {:?}", e)
@@ -441,14 +439,11 @@ impl JobQueue {
                     .get()
                     .ok_or_else(|| "Workflow engine not initialized".to_string())?;
 
-                let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
+                let workflow_data =
-                // Convert Box to Arc for context storage
+                    create_wasm_removal_workflow_data(*request.clone(), Arc::clone(context));
-                let request_arc: Arc<WasmModuleRemovalRequest> = Arc::new(*request.clone());
-                workflow_context.set_arc("wasm_module_removal_request", request_arc);
-                workflow_context.set_arc("app_context", Arc::clone(context));
 
                 let instance_id = engine
-                    .start_workflow(WorkflowId::new("wasm_module_removal"), workflow_context)
+                    .start_workflow(WorkflowId::new("wasm_module_removal"), workflow_data)
                     .await
                     .map_err(|e| {
                         format!("Failed to start WASM module removal workflow: {:?}", e)
@@ -674,13 +669,11 @@ impl JobQueue {
                     .get()
                     .ok_or_else(|| "Workflow engine not initialized".to_string())?;
 
-                let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
+                let workflow_data =
-                let config_arc: Arc<TokenizerConfigRequest> = Arc::new(*config.clone());
+                    create_tokenizer_workflow_data(*config.clone(), Arc::clone(context));
-                workflow_context.set_arc("tokenizer_config", config_arc);
-                workflow_context.set_arc("app_context", Arc::clone(context));
 
                 let instance_id = engine
-                    .start_workflow(WorkflowId::new("tokenizer_registration"), workflow_context)
+                    .start_workflow(WorkflowId::new("tokenizer_registration"), workflow_data)
                     .await
                     .map_err(|e| {
                         format!("Failed to start tokenizer registration workflow: {:?}", e)
@@ -719,86 +712,82 @@ impl JobQueue {
 
     /// Start a workflow and return its instance ID
     async fn start_worker_workflow(
-        engine: &Arc<WorkflowEngine>,
+        engine: &Arc<AppWorkflowEngine>,
         config: &WorkerConfigRequest,
         context: &Arc<AppContext>,
     ) -> Result<WorkflowInstanceId, String> {
-        let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
-        workflow_context.set("worker_config", config.clone());
-        workflow_context.set_arc("app_context", Arc::clone(context));
-
         // Select workflow based on runtime field
-        let workflow_id = match config.runtime.as_deref() {
+        let (workflow_id, workflow_data) = match config.runtime.as_deref() {
-            Some("external") => WorkflowId::new("external_worker_registration"),
+            Some("external") => (
-            _ => WorkflowId::new("local_worker_registration"),
+                WorkflowId::new("external_worker_registration"),
+                create_external_worker_workflow_data(config.clone(), Arc::clone(context)),
+            ),
+            _ => (
+                WorkflowId::new("local_worker_registration"),
+                create_local_worker_workflow_data(config.clone(), Arc::clone(context)),
+            ),
         };
 
         engine
-            .start_workflow(workflow_id, workflow_context)
+            .start_workflow(workflow_id, workflow_data)
             .await
             .map_err(|e| format!("Failed to start worker registration workflow: {:?}", e))
     }
 
     /// Start worker removal workflow
     async fn start_worker_removal_workflow(
-        engine: &Arc<WorkflowEngine>,
+        engine: &Arc<AppWorkflowEngine>,
         url: &str,
         context: &Arc<AppContext>,
     ) -> Result<WorkflowInstanceId, String> {
-        let removal_request = WorkerRemovalRequest {
+        let workflow_data = create_worker_removal_workflow_data(
-            url: url.to_string(),
+            url.to_string(),
-            dp_aware: context.router_config.dp_aware,
+            context.router_config.dp_aware,
-        };
+            Arc::clone(context),
-
+        );
-        let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
-        workflow_context.set("removal_request", removal_request);
-        workflow_context.set_arc("app_context", Arc::clone(context));
 
         engine
-            .start_workflow(WorkflowId::new("worker_removal"), workflow_context)
+            .start_workflow(WorkflowId::new("worker_removal"), workflow_data)
             .await
             .map_err(|e| format!("Failed to start worker removal workflow: {:?}", e))
     }
 
     /// Start worker update workflow
     async fn start_worker_update_workflow(
-        engine: &Arc<WorkflowEngine>,
+        engine: &Arc<AppWorkflowEngine>,
         url: &str,
         update: &WorkerUpdateRequest,
         context: &Arc<AppContext>,
     ) -> Result<WorkflowInstanceId, String> {
-        let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
+        let workflow_data = create_worker_update_workflow_data(
-        // Pass URL and dp_aware separately, workflow step handles the rest
+            url.to_string(),
-        workflow_context.set("worker_url", url.to_string());
+            update.clone(),
-        workflow_context.set("dp_aware", context.router_config.dp_aware);
+            Arc::clone(context),
-        workflow_context.set("update_request", update.clone());
+        );
-        workflow_context.set_arc("app_context", Arc::clone(context));
 
         engine
-            .start_workflow(WorkflowId::new("worker_update"), workflow_context)
+            .start_workflow(WorkflowId::new("worker_update"), workflow_data)
             .await
             .map_err(|e| format!("Failed to start worker update workflow: {:?}", e))
     }
 
     /// Start MCP server registration workflow
     async fn start_mcp_registration_workflow(
-        engine: &Arc<WorkflowEngine>,
+        engine: &Arc<AppWorkflowEngine>,
         config: &McpServerConfigRequest,
         context: &Arc<AppContext>,
     ) -> Result<WorkflowInstanceId, String> {
-        let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
+        let workflow_data = create_mcp_workflow_data(config.clone(), Arc::clone(context));
-        workflow_context.set("mcp_server_config", config.clone());
-        workflow_context.set_arc("app_context", Arc::clone(context));
 
         engine
-            .start_workflow(WorkflowId::new("mcp_registration"), workflow_context)
+            .start_workflow(WorkflowId::new("mcp_registration"), workflow_data)
             .await
             .map_err(|e| format!("Failed to start MCP registration workflow: {:?}", e))
     }
 
     /// Wait for workflow completion with adaptive polling
     async fn wait_for_workflow_completion(
-        engine: &Arc<WorkflowEngine>,
+        engine: &Arc<AppWorkflowEngine>,
         instance_id: WorkflowInstanceId,
         worker_url: &str,
         timeout_duration: Duration,

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

@@ -1,18 +1,21 @@
 use std::{sync::Arc, time::Duration};
 
 use async_trait::async_trait;
-use rmcp::{service::RunningService, RoleClient};
 use tracing::{debug, error, info, warn};
 
+use super::workflow_data::{AnyWorkflowData, McpWorkflowData};
 use crate::{
     app_context::AppContext,
     mcp::{config::McpServerConfig, manager::McpManager},
     observability::metrics::Metrics,
-    workflow::*,
+    workflow::{
+        BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId,
+        StepResult, WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult,
+    },
 };
 
 /// MCP server connection configuration
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct McpServerConfigRequest {
     /// Server name (unique identifier)
     pub name: String,
@@ -35,11 +38,17 @@ impl McpServerConfigRequest {
 pub struct ConnectMcpServerStep;
 
 #[async_trait]
-impl StepExecutor for ConnectMcpServerStep {
+impl StepExecutor<AnyWorkflowData> for ConnectMcpServerStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config_request: Arc<McpServerConfigRequest> =
+        &self,
-            context.get_or_err("mcp_server_config")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_mcp()?;
+        let config_request = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
 
         debug!("Connecting to MCP server: {}", config_request.name);
 
@@ -66,8 +75,9 @@ impl StepExecutor for ConnectMcpServerStep {
             config_request.name
         );
 
-        // Store client in context (context.set() will wrap in Arc)
+        // Store client in typed data
-        context.set("mcp_client", client);
+        let data_mut = context.data.as_mcp_mut()?;
+        data_mut.mcp_client = Some(Arc::new(client));
 
         Ok(StepResult::Success)
     }
@@ -86,12 +96,21 @@ impl StepExecutor for ConnectMcpServerStep {
 pub struct DiscoverMcpInventoryStep;
 
 #[async_trait]
-impl StepExecutor for DiscoverMcpInventoryStep {
+impl StepExecutor<AnyWorkflowData> for DiscoverMcpInventoryStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config_request: Arc<McpServerConfigRequest> =
+        &self,
-            context.get_or_err("mcp_server_config")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+    ) -> WorkflowResult<StepResult> {
-        let mcp_client: Arc<RunningService<RoleClient, ()>> = context.get_or_err("mcp_client")?;
+        let data = context.data.as_mcp()?;
+        let config_request = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let mcp_client = data
+            .mcp_client
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?;
 
         debug!(
             "Discovering inventory for MCP server: {}",
@@ -111,7 +130,7 @@ impl StepExecutor for DiscoverMcpInventoryStep {
         let inventory = mcp_manager.inventory();
 
         // Use the public load_server_inventory method
-        McpManager::load_server_inventory(&inventory, &config_request.name, &mcp_client).await;
+        McpManager::load_server_inventory(&inventory, &config_request.name, mcp_client).await;
 
         info!("Completed inventory discovery for {}", config_request.name);
 
@@ -130,12 +149,22 @@ impl StepExecutor for DiscoverMcpInventoryStep {
 pub struct RegisterMcpServerStep;
 
 #[async_trait]
-impl StepExecutor for RegisterMcpServerStep {
+impl StepExecutor<AnyWorkflowData> for RegisterMcpServerStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config_request: Arc<McpServerConfigRequest> =
+        &self,
-            context.get_or_err("mcp_server_config")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+    ) -> WorkflowResult<StepResult> {
-        let mcp_client: Arc<RunningService<RoleClient, ()>> = context.get_or_err("mcp_client")?;
+        let data = context.data.as_mcp()?;
+        let config_request = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let mcp_client = data
+            .mcp_client
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?
+            .clone();
 
         debug!("Registering MCP server: {}", config_request.name);
 
@@ -174,20 +203,26 @@ impl StepExecutor for RegisterMcpServerStep {
 pub struct ValidateRegistrationStep;
 
 #[async_trait]
-impl StepExecutor for ValidateRegistrationStep {
+impl StepExecutor<AnyWorkflowData> for ValidateRegistrationStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config_request: Arc<McpServerConfigRequest> =
+        &self,
-            context.get_or_err("mcp_server_config")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_mcp()?;
+        let config_request = &data.config;
 
-        let client_registered = context
+        let client_registered = data.mcp_client.is_some();
-            .get::<RunningService<RoleClient, ()>>("mcp_client")
-            .is_some();
 
         if client_registered {
             info!(
                 "MCP server '{}' registered successfully",
                 config_request.name
             );
+
+            // Mark as validated
+            let data_mut = context.data.as_mcp_mut()?;
+            data_mut.validated = true;
+
             return Ok(StepResult::Success);
         }
 
@@ -228,7 +263,7 @@ impl StepExecutor for ValidateRegistrationStep {
 /// - DiscoverMcpInventory: 3 retries, 10s timeout (discovery + caching)
 /// - RegisterMcpServer: No retry, 5s timeout (fast registration)
 /// - ValidateRegistration: Final validation step
-pub fn create_mcp_registration_workflow() -> WorkflowDefinition {
+pub fn create_mcp_registration_workflow() -> WorkflowDefinition<AnyWorkflowData> {
     WorkflowDefinition::new("mcp_registration", "MCP Server Registration")
         .add_step(
             StepDefinition::new(
@@ -281,3 +316,16 @@ pub fn create_mcp_registration_workflow() -> WorkflowDefinition {
             .depends_on(&["register_mcp_server"]),
         )
 }
+
+/// Helper to create initial workflow data for MCP registration
+pub fn create_mcp_workflow_data(
+    config: McpServerConfigRequest,
+    app_context: Arc<AppContext>,
+) -> AnyWorkflowData {
+    AnyWorkflowData::Mcp(McpWorkflowData {
+        config,
+        validated: false,
+        app_context: Some(app_context),
+        mcp_client: None,
+    })
+}

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

@@ -11,19 +11,38 @@ pub mod tokenizer_registration;
 pub mod wasm_module_registration;
 pub mod wasm_module_removal;
 pub mod worker;
+pub mod workflow_data;
 
 // Worker management (registration, removal)
-#[allow(deprecated)]
+pub use mcp_registration::{
-pub use worker::create_external_worker_registration_workflow;
+    create_mcp_registration_workflow, create_mcp_workflow_data, ConnectMcpServerStep,
-// Backward compatibility aliases
+    DiscoverMcpInventoryStep, McpServerConfigRequest, RegisterMcpServerStep,
-#[allow(deprecated)]
+    ValidateRegistrationStep,
-pub use worker::create_worker_registration_workflow;
+};
+pub use tokenizer_registration::{
+    create_tokenizer_registration_workflow, create_tokenizer_workflow_data, LoadTokenizerStep,
+    TokenizerConfigRequest, TokenizerRemovalRequest, ValidateTokenizerConfigStep,
+};
+pub use wasm_module_registration::{
+    create_wasm_module_registration_workflow, create_wasm_registration_workflow_data,
+    CalculateHashStep, CheckDuplicateStep, LoadWasmBytesStep, RegisterModuleStep,
+    ValidateDescriptorStep, ValidateWasmComponentStep, WasmModuleConfigRequest,
+};
+pub use wasm_module_removal::{
+    create_wasm_module_removal_workflow, create_wasm_removal_workflow_data, FindModuleToRemoveStep,
+    RemoveModuleStep, WasmModuleRemovalRequest,
+};
 pub use worker::{
     // Workflow builders
     create_external_worker_workflow,
+    // Workflow data helpers
+    create_external_worker_workflow_data,
     create_local_worker_workflow,
+    create_local_worker_workflow_data,
     create_worker_removal_workflow,
+    create_worker_removal_workflow_data,
     create_worker_update_workflow,
+    create_worker_update_workflow_data,
     // Utility functions
     group_models_into_cards,
     infer_model_type_from_id,
@@ -54,29 +73,10 @@ pub use worker::{
     WorkerList,
     WorkerRemovalRequest,
 };
-
+// Typed workflow data structures
-// Legacy type aliases for backward compatibility
+pub use workflow_data::{
-pub type ActivateWorkerStep = ActivateWorkersStep;
+    AnyWorkflowData, ExternalWorkerWorkflowData, LocalWorkerWorkflowData, McpWorkflowData,
-pub type RegisterWorkerStep = RegisterWorkersStep;
+    ProtocolUpdateRequest, TokenizerWorkflowData, WasmRegistrationWorkflowData,
-pub type CreateWorkerStep = CreateLocalWorkerStep;
+    WasmRemovalWorkflowData, WorkerConfigRequest, WorkerList as WorkflowWorkerList,
-pub type ActivateExternalWorkersStep = ActivateWorkersStep;
+    WorkerRemovalWorkflowData, WorkerUpdateWorkflowData,
-pub type RegisterExternalWorkersStep = RegisterWorkersStep;
-pub type UpdateExternalPoliciesStep = UpdatePoliciesStep;
-
-pub use mcp_registration::{
-    create_mcp_registration_workflow, ConnectMcpServerStep, DiscoverMcpInventoryStep,
-    McpServerConfigRequest, RegisterMcpServerStep, ValidateRegistrationStep,
-};
-pub use tokenizer_registration::{
-    create_tokenizer_registration_workflow, LoadTokenizerStep, TokenizerConfigRequest,
-    TokenizerRemovalRequest, ValidateTokenizerConfigStep,
-};
-pub use wasm_module_registration::{
-    create_wasm_module_registration_workflow, CalculateHashStep, CheckDuplicateStep,
-    LoadWasmBytesStep, RegisterModuleStep, ValidateDescriptorStep, ValidateWasmComponentStep,
-    WasmModuleConfigRequest,
-};
-pub use wasm_module_removal::{
-    create_wasm_module_removal_workflow, FindModuleToRemoveStep, RemoveModuleStep,
-    WasmModuleRemovalRequest,
 };

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

@@ -9,7 +9,15 @@ use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use tracing::{debug, error, info};
 
-use crate::{app_context::AppContext, tokenizer::factory, workflow::*};
+use super::workflow_data::{AnyWorkflowData, TokenizerWorkflowData};
+use crate::{
+    app_context::AppContext,
+    tokenizer::factory,
+    workflow::{
+        BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId,
+        StepResult, WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult,
+    },
+};
 
 /// Configuration for adding a tokenizer
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -39,10 +47,17 @@ pub struct TokenizerRemovalRequest {
 pub struct ValidateTokenizerConfigStep;
 
 #[async_trait]
-impl StepExecutor for ValidateTokenizerConfigStep {
+impl StepExecutor<AnyWorkflowData> for ValidateTokenizerConfigStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config: Arc<TokenizerConfigRequest> = context.get_or_err("tokenizer_config")?;
+        &self,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_tokenizer()?;
+        let config = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
 
         debug!(
             "Validating tokenizer config: name={}, source={}",
@@ -86,22 +101,36 @@ impl StepExecutor for ValidateTokenizerConfigStep {
 pub struct LoadTokenizerStep;
 
 #[async_trait]
-impl StepExecutor for LoadTokenizerStep {
+impl StepExecutor<AnyWorkflowData> for LoadTokenizerStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config: Arc<TokenizerConfigRequest> = context.get_or_err("tokenizer_config")?;
+        &self,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_tokenizer()?;
+        let config = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
+            .clone();
 
         info!(
             "Loading tokenizer '{}' (id: {}) from source: {}",
             config.name, config.id, config.source
         );
 
+        // Clone needed values before async move
+        let id = config.id.clone();
+        let name = config.name.clone();
+        let source = config.source.clone();
+        let chat_template = config.chat_template_path.clone();
+
         // Load the tokenizer using the registry's load method (handles deduplication)
         let result = app_context
             .tokenizer_registry
-            .load(&config.id, &config.name, &config.source, || {
+            .load(&id, &name, &source, || {
-                let source = config.source.clone();
+                let source = source.clone();
-                let chat_template = config.chat_template_path.clone();
+                let chat_template = chat_template.clone();
                 async move {
                     factory::create_tokenizer_async_with_chat_template(
                         &source,
@@ -123,18 +152,19 @@ impl StepExecutor for LoadTokenizerStep {
 
                 info!(
                     "Successfully loaded tokenizer '{}' (id: {}) with vocab_size: {:?}",
-                    config.name, loaded_id, vocab_size
+                    name, loaded_id, vocab_size
                 );
 
-                // Store vocab size in context for later use
+                // Store vocab size in typed data
                 if let Some(size) = vocab_size {
-                    context.set("vocab_size", size);
+                    let data_mut = context.data.as_tokenizer_mut()?;
+                    data_mut.vocab_size = Some(size);
                 }
 
                 Ok(StepResult::Success)
             }
             Err(e) => {
-                error!("Failed to load tokenizer '{}': {}", config.name, e);
+                error!("Failed to load tokenizer '{}': {}", name, e);
                 Err(WorkflowError::StepFailed {
                     step_id: StepId::new("load_tokenizer"),
                     message: e,
@@ -161,7 +191,7 @@ impl StepExecutor for LoadTokenizerStep {
 /// Workflow configuration:
 /// - ValidateConfig: No retry, 5s timeout (fast validation)
 /// - LoadTokenizer: 3 retries, 5min timeout (may need to download from HuggingFace)
-pub fn create_tokenizer_registration_workflow() -> WorkflowDefinition {
+pub fn create_tokenizer_registration_workflow() -> WorkflowDefinition<AnyWorkflowData> {
     WorkflowDefinition::new("tokenizer_registration", "Tokenizer Registration")
         .add_step(
             StepDefinition::new(
@@ -188,6 +218,18 @@ pub fn create_tokenizer_registration_workflow() -> WorkflowDefinition {
         )
 }
 
+/// Helper to create initial workflow data for tokenizer registration
+pub fn create_tokenizer_workflow_data(
+    config: TokenizerConfigRequest,
+    app_context: Arc<AppContext>,
+) -> AnyWorkflowData {
+    AnyWorkflowData::Tokenizer(TokenizerWorkflowData {
+        config,
+        vocab_size: None,
+        app_context: Some(app_context),
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -212,7 +254,11 @@ mod tests {
 
     #[test]
     fn test_workflow_creation() {
-        let workflow = create_tokenizer_registration_workflow();
+        let mut workflow = create_tokenizer_registration_workflow();
         assert_eq!(workflow.id.to_string(), "tokenizer_registration");
+        // Validate the workflow DAG
+        workflow
+            .validate()
+            .expect("Workflow validation should pass");
     }
 }

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

@@ -10,14 +10,18 @@ use tracing::{debug, info, warn};
 use uuid::Uuid;
 use wasmtime::{component::Component, Config, Engine};
 
+use super::workflow_data::{AnyWorkflowData, WasmRegistrationWorkflowData};
 use crate::{
     app_context::AppContext,
     wasm::module::{WasmModule, WasmModuleDescriptor, WasmModuleMeta},
-    workflow::*,
+    workflow::{
+        BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId,
+        StepResult, WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult,
+    },
 };
 
 /// WASM module registration request
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct WasmModuleConfigRequest {
     /// Module descriptor containing name, file_path, attach_points, etc.
     pub descriptor: WasmModuleDescriptor,
@@ -61,12 +65,13 @@ fn has_wasm_extension(path: &Path) -> bool {
 pub struct ValidateDescriptorStep;
 
 #[async_trait]
-impl StepExecutor for ValidateDescriptorStep {
+impl StepExecutor<AnyWorkflowData> for ValidateDescriptorStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config_request: Arc<WasmModuleConfigRequest> =
+        &self,
-            context.get_or_err("wasm_module_config")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-
+    ) -> WorkflowResult<StepResult> {
-        let descriptor = &config_request.descriptor;
+        let data = context.data.as_wasm_registration()?;
+        let descriptor = &data.config.descriptor;
 
         debug!("Validating WASM module descriptor: {}", descriptor.name);
 
@@ -198,12 +203,16 @@ impl StepExecutor for ValidateDescriptorStep {
             });
         }
 
-        // Store file size in context for later steps
+        // Clone name for logging before mutable borrow
-        context.set("file_size_bytes", metadata.len());
+        let module_name = descriptor.name.clone();
+
+        // Store file size in typed data
+        let data_mut = context.data.as_wasm_registration_mut()?;
+        data_mut.file_size_bytes = Some(metadata.len());
 
         info!(
             "Descriptor validated successfully for module: {}",
-            descriptor.name
+            module_name
         );
         Ok(StepResult::Success)
     }
@@ -220,12 +229,13 @@ impl StepExecutor for ValidateDescriptorStep {
 pub struct CalculateHashStep;
 
 #[async_trait]
-impl StepExecutor for CalculateHashStep {
+impl StepExecutor<AnyWorkflowData> for CalculateHashStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config_request: Arc<WasmModuleConfigRequest> =
+        &self,
-            context.get_or_err("wasm_module_config")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-
+    ) -> WorkflowResult<StepResult> {
-        let file_path = &config_request.descriptor.file_path;
+        let data = context.data.as_wasm_registration()?;
+        let file_path = &data.config.descriptor.file_path;
 
         debug!("Calculating SHA256 hash for: {}", file_path);
 
@@ -260,10 +270,14 @@ impl StepExecutor for CalculateHashStep {
 
         let hash: [u8; 32] = hasher.finalize().into();
 
-        // Store hash in context
+        // Clone path for logging before mutable borrow
-        context.set("sha256_hash", hash);
+        let path_for_log = file_path.clone();
 
-        info!("SHA256 hash calculated for: {}", file_path);
+        // Store hash in typed data
+        let data_mut = context.data.as_wasm_registration_mut()?;
+        data_mut.sha256_hash = Some(hash);
+
+        info!("SHA256 hash calculated for: {}", path_for_log);
         Ok(StepResult::Success)
     }
 
@@ -279,16 +293,24 @@ impl StepExecutor for CalculateHashStep {
 pub struct CheckDuplicateStep;
 
 #[async_trait]
-impl StepExecutor for CheckDuplicateStep {
+impl StepExecutor<AnyWorkflowData> for CheckDuplicateStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config_request: Arc<WasmModuleConfigRequest> =
+        &self,
-            context.get_or_err("wasm_module_config")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+    ) -> WorkflowResult<StepResult> {
-        let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?;
+        let data = context.data.as_wasm_registration()?;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let sha256_hash = data
+            .sha256_hash
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("sha256_hash".to_string()))?;
 
         debug!(
             "Checking for duplicate SHA256 hash for module: {}",
-            config_request.descriptor.name
+            data.config.descriptor.name
         );
 
         // Get WASM module manager from app context
@@ -303,7 +325,7 @@ impl StepExecutor for CheckDuplicateStep {
 
         // Check for duplicate hash using manager's internal method
         wasm_manager
-            .check_duplicate_sha256_hash(sha256_hash.as_ref())
+            .check_duplicate_sha256_hash(sha256_hash)
             .map_err(|e| WorkflowError::StepFailed {
                 step_id: StepId::new("check_duplicate"),
                 message: format!("Duplicate SHA256 hash detected: {}", e),
@@ -311,7 +333,7 @@ impl StepExecutor for CheckDuplicateStep {
 
         info!(
             "No duplicate found for module: {}",
-            config_request.descriptor.name
+            data.config.descriptor.name
         );
         Ok(StepResult::Success)
     }
@@ -328,15 +350,19 @@ impl StepExecutor for CheckDuplicateStep {
 pub struct LoadWasmBytesStep;
 
 #[async_trait]
-impl StepExecutor for LoadWasmBytesStep {
+impl StepExecutor<AnyWorkflowData> for LoadWasmBytesStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config_request: Arc<WasmModuleConfigRequest> =
+        &self,
-            context.get_or_err("wasm_module_config")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-
+    ) -> WorkflowResult<StepResult> {
-        let file_path = &config_request.descriptor.file_path;
+        let data = context.data.as_wasm_registration()?;
+        let file_path = &data.config.descriptor.file_path;
 
         debug!("Loading WASM bytes from: {}", file_path);
 
+        // Clone path for logging before mutable borrow
+        let path_for_log = file_path.clone();
+
         let wasm_bytes =
             tokio::fs::read(file_path)
                 .await
@@ -345,10 +371,11 @@ impl StepExecutor for LoadWasmBytesStep {
                     message: format!("Failed to read WASM file {}: {}", file_path, e),
                 })?;
 
-        // Store WASM bytes in context
+        // Store WASM bytes in typed data
-        context.set("wasm_bytes", wasm_bytes);
+        let data_mut = context.data.as_wasm_registration_mut()?;
+        data_mut.wasm_bytes = Some(wasm_bytes);
 
-        info!("WASM bytes loaded from: {}", file_path);
+        info!("WASM bytes loaded from: {}", path_for_log);
         Ok(StepResult::Success)
     }
 
@@ -364,15 +391,20 @@ impl StepExecutor for LoadWasmBytesStep {
 pub struct ValidateWasmComponentStep;
 
 #[async_trait]
-impl StepExecutor for ValidateWasmComponentStep {
+impl StepExecutor<AnyWorkflowData> for ValidateWasmComponentStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config_request: Arc<WasmModuleConfigRequest> =
+        &self,
-            context.get_or_err("wasm_module_config")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let wasm_bytes: Arc<Vec<u8>> = context.get_or_err("wasm_bytes")?;
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_wasm_registration()?;
+        let wasm_bytes = data
+            .wasm_bytes
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_bytes".to_string()))?;
 
         debug!(
             "Validating WASM component format for module: {}",
-            config_request.descriptor.name
+            data.config.descriptor.name
         );
 
         // Create a temporary engine to validate the component
@@ -386,7 +418,7 @@ impl StepExecutor for ValidateWasmComponentStep {
         })?;
 
         // Attempt to compile the component to validate it
-        Component::new(&engine, wasm_bytes.as_ref())
+        Component::new(&engine, wasm_bytes)
             .map_err(|e| WorkflowError::StepFailed {
                 step_id: StepId::new("validate_wasm_component"),
                 message: format!(
@@ -399,7 +431,7 @@ impl StepExecutor for ValidateWasmComponentStep {
 
         info!(
             "WASM component validated successfully for module: {}",
-            config_request.descriptor.name
+            data.config.descriptor.name
         );
         Ok(StepResult::Success)
     }
@@ -416,19 +448,31 @@ impl StepExecutor for ValidateWasmComponentStep {
 pub struct RegisterModuleStep;
 
 #[async_trait]
-impl StepExecutor for RegisterModuleStep {
+impl StepExecutor<AnyWorkflowData> for RegisterModuleStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config_request: Arc<WasmModuleConfigRequest> =
+        &self,
-            context.get_or_err("wasm_module_config")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+    ) -> WorkflowResult<StepResult> {
-        let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?;
+        let data = context.data.as_wasm_registration()?;
-        let file_size_bytes: Arc<u64> = context.get_or_err("file_size_bytes")?;
+        let app_context = data
-        let wasm_bytes: Arc<Vec<u8>> = context.get_or_err("wasm_bytes")?;
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let sha256_hash = data
+            .sha256_hash
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("sha256_hash".to_string()))?;
+        let file_size_bytes = data
+            .file_size_bytes
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("file_size_bytes".to_string()))?;
+        let wasm_bytes = data
+            .wasm_bytes
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_bytes".to_string()))?
+            .clone();
 
-        debug!(
+        let descriptor = &data.config.descriptor;
-            "Registering WASM module in manager: {}",
+
-            config_request.descriptor.name
+        debug!("Registering WASM module in manager: {}", descriptor.name);
-        );
 
         // Get WASM module manager from app context
         let wasm_manager =
@@ -451,18 +495,21 @@ impl StepExecutor for RegisterModuleStep {
         let module = WasmModule {
             module_uuid,
             module_meta: WasmModuleMeta {
-                name: config_request.descriptor.name.clone(),
+                name: descriptor.name.clone(),
-                file_path: config_request.descriptor.file_path.clone(),
+                file_path: descriptor.file_path.clone(),
-                sha256_hash: *sha256_hash.as_ref(),
+                sha256_hash,
-                size_bytes: *file_size_bytes.as_ref(),
+                size_bytes: file_size_bytes,
                 created_at: now,
                 last_accessed_at: now,
                 access_count: 0,
-                attach_points: config_request.descriptor.attach_points.clone(),
+                attach_points: descriptor.attach_points.clone(),
-                wasm_bytes: wasm_bytes.as_ref().clone(),
+                wasm_bytes,
             },
         };
 
+        // Clone name for logging before mutable borrow
+        let module_name = descriptor.name.clone();
+
         // Register module in manager
         wasm_manager
             .register_module_internal(module)
@@ -471,12 +518,13 @@ impl StepExecutor for RegisterModuleStep {
                 message: format!("Failed to register module: {}", e),
             })?;
 
-        // Store module UUID in context for return value
+        // Store module UUID in typed data
-        context.set("module_uuid", module_uuid);
+        let data_mut = context.data.as_wasm_registration_mut()?;
+        data_mut.module_uuid = Some(module_uuid);
 
         info!(
             "WASM module registered successfully: {} (UUID: {})",
-            config_request.descriptor.name, module_uuid
+            module_name, module_uuid
         );
 
         Ok(StepResult::Success)
@@ -504,7 +552,7 @@ impl StepExecutor for RegisterModuleStep {
 /// - LoadWasmBytes: 3 retries, 60s timeout (I/O intensive)
 /// - ValidateWasmComponent: No retry, 30s timeout (CPU intensive validation)
 /// - RegisterModule: No retry, 5s timeout (fast registration)
-pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition {
+pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition<AnyWorkflowData> {
     WorkflowDefinition::new("wasm_module_registration", "WASM Module Registration")
         .add_step(
             StepDefinition::new(
@@ -574,3 +622,18 @@ pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition {
             .depends_on(&["validate_wasm_component"]),
         )
 }
+
+/// Helper to create initial workflow data for WASM module registration
+pub fn create_wasm_registration_workflow_data(
+    config: WasmModuleConfigRequest,
+    app_context: Arc<AppContext>,
+) -> AnyWorkflowData {
+    AnyWorkflowData::WasmRegistration(WasmRegistrationWorkflowData {
+        config,
+        wasm_bytes: None,
+        sha256_hash: None,
+        file_size_bytes: None,
+        module_uuid: None,
+        app_context: Some(app_context),
+    })
+}

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

@@ -4,10 +4,17 @@ use async_trait::async_trait;
 use tracing::{debug, info};
 use uuid::Uuid;
 
-use crate::{app_context::AppContext, workflow::*};
+use super::workflow_data::{AnyWorkflowData, WasmRemovalWorkflowData};
+use crate::{
+    app_context::AppContext,
+    workflow::{
+        FailureAction, StepDefinition, StepExecutor, StepId, StepResult, WorkflowContext,
+        WorkflowDefinition, WorkflowError, WorkflowResult,
+    },
+};
 
 /// WASM module removal request
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct WasmModuleRemovalRequest {
     /// Module UUID to remove
     pub module_uuid: Uuid,
@@ -30,11 +37,17 @@ impl WasmModuleRemovalRequest {
 pub struct FindModuleToRemoveStep;
 
 #[async_trait]
-impl StepExecutor for FindModuleToRemoveStep {
+impl StepExecutor<AnyWorkflowData> for FindModuleToRemoveStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let removal_request: Arc<WasmModuleRemovalRequest> =
+        &self,
-            context.get_or_err("wasm_module_removal_request")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_wasm_removal()?;
+        let removal_request = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
 
         debug!("Finding module to remove: {}", removal_request.module_uuid);
 
@@ -63,7 +76,14 @@ impl StepExecutor for FindModuleToRemoveStep {
             });
         }
 
-        info!("Module found for removal: {}", removal_request.module_uuid);
+        // Clone uuid for logging before mutable borrow
+        let module_uuid = removal_request.module_uuid;
+
+        // Store the module ID in typed data
+        let data_mut = context.data.as_wasm_removal_mut()?;
+        data_mut.module_id = Some(module_uuid.to_string());
+
+        info!("Module found for removal: {}", module_uuid);
         Ok(StepResult::Success)
     }
 
@@ -78,11 +98,17 @@ impl StepExecutor for FindModuleToRemoveStep {
 pub struct RemoveModuleStep;
 
 #[async_trait]
-impl StepExecutor for RemoveModuleStep {
+impl StepExecutor<AnyWorkflowData> for RemoveModuleStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let removal_request: Arc<WasmModuleRemovalRequest> =
+        &self,
-            context.get_or_err("wasm_module_removal_request")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_wasm_removal()?;
+        let removal_request = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
 
         debug!("Removing WASM module: {}", removal_request.module_uuid);
 
@@ -125,7 +151,7 @@ impl StepExecutor for RemoveModuleStep {
 /// Workflow configuration:
 /// - FindModuleToRemove: No retry, 5s timeout (fast lookup)
 /// - RemoveModule: No retry, 5s timeout (fast removal)
-pub fn create_wasm_module_removal_workflow() -> WorkflowDefinition {
+pub fn create_wasm_module_removal_workflow() -> WorkflowDefinition<AnyWorkflowData> {
     WorkflowDefinition::new("wasm_module_removal", "WASM Module Removal")
         .add_step(
             StepDefinition::new(
@@ -143,3 +169,15 @@ pub fn create_wasm_module_removal_workflow() -> WorkflowDefinition {
                 .depends_on(&["find_module_to_remove"]),
         )
 }
+
+/// Helper to create initial workflow data for WASM module removal
+pub fn create_wasm_removal_workflow_data(
+    config: WasmModuleRemovalRequest,
+    app_context: Arc<AppContext>,
+) -> AnyWorkflowData {
+    AnyWorkflowData::WasmRemoval(WasmRemovalWorkflowData {
+        config,
+        module_id: None,
+        app_context: Some(app_context),
+    })
+}

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

@@ -6,14 +6,12 @@ use async_trait::async_trait;
 use tracing::{debug, info};
 
 use crate::{
-    app_context::AppContext,
     core::{
         circuit_breaker::CircuitBreakerConfig,
-        model_card::ModelCard,
+        steps::workflow_data::{AnyWorkflowData, WorkerList},
         worker::{HealthConfig, RuntimeType, WorkerType},
         BasicWorkerBuilder, ConnectionMode, Worker,
     },
-    protocols::worker_spec::WorkerConfigRequest,
     workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 
@@ -30,11 +28,18 @@ fn normalize_external_url(url: &str) -> String {
 pub struct CreateExternalWorkersStep;
 
 #[async_trait]
-impl StepExecutor for CreateExternalWorkersStep {
+impl StepExecutor<AnyWorkflowData> for CreateExternalWorkersStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
+        &self,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let model_cards: Arc<Vec<ModelCard>> = context.get_or_err("model_cards")?;
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_external_worker()?;
+        let config = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let model_cards = &data.model_cards;
 
         // Build configs from router settings
         let circuit_breaker_config = {
@@ -144,8 +149,11 @@ impl StepExecutor for CreateExternalWorkersStep {
             );
         }
 
-        context.set("workers", workers);
+        // Store results in workflow data
-        context.set("labels", labels);
+        let data_mut = context.data.as_external_worker_mut()?;
+        data_mut.workers = Some(WorkerList::from_workers(&workers));
+        data_mut.actual_workers = Some(workers);
+        data_mut.labels = labels;
         Ok(StepResult::Success)
     }
 

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

@@ -1,6 +1,6 @@
 //! Model discovery step for external API endpoints.
 
-use std::{collections::HashMap, sync::Arc, time::Duration};
+use std::{collections::HashMap, time::Duration};
 
 use async_trait::async_trait;
 use once_cell::sync::Lazy;
@@ -13,8 +13,8 @@ use crate::{
     core::{
         model_card::{ModelCard, ProviderType},
         model_type::ModelType,
+        steps::workflow_data::AnyWorkflowData,
     },
-    protocols::worker_spec::WorkerConfigRequest,
     workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 
@@ -225,9 +225,13 @@ async fn fetch_models(url: &str, api_key: Option<&str>) -> Result<Vec<ModelCard>
 pub struct DiscoverModelsStep;
 
 #[async_trait]
-impl StepExecutor for DiscoverModelsStep {
+impl StepExecutor<AnyWorkflowData> for DiscoverModelsStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
+        &self,
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_external_worker()?;
+        let config = &data.config;
 
         // If no API key is provided, skip model discovery and use wildcard mode.
         if config.api_key.as_ref().is_none_or(|k| k.is_empty()) {
@@ -236,7 +240,7 @@ impl StepExecutor for DiscoverModelsStep {
                  User's Authorization header will be forwarded to backend.",
                 config.url
             );
-            context.set::<Vec<ModelCard>>("model_cards", vec![]);
+            // Leave model_cards empty for wildcard mode
             return Ok(StepResult::Success);
         }
 
@@ -263,7 +267,7 @@ impl StepExecutor for DiscoverModelsStep {
             model_cards.iter().map(|c| &c.id).collect::<Vec<_>>()
         );
 
-        context.set("model_cards", model_cards);
+        context.data.as_external_worker_mut()?.model_cards = model_cards;
         Ok(StepResult::Success)
     }
 

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

@@ -15,8 +15,11 @@ pub use discover_models::{
 };
 
 use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep};
-use crate::workflow::{
+use crate::{
-    BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition,
+    app_context::AppContext,
+    core::steps::workflow_data::{AnyWorkflowData, ExternalWorkerWorkflowData},
+    protocols::worker_spec::WorkerConfigRequest,
+    workflow::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition},
 };
 
 /// Create external worker registration workflow definition.
@@ -35,7 +38,7 @@ use crate::workflow::{
 ///       │                         │
 ///       └────────────┴────────────┘
 /// ```
-pub fn create_external_worker_workflow() -> WorkflowDefinition {
+pub fn create_external_worker_workflow() -> WorkflowDefinition<AnyWorkflowData> {
     WorkflowDefinition::new(
         "external_worker_registration",
         "External Worker Registration",
@@ -102,3 +105,18 @@ pub fn create_external_worker_workflow() -> WorkflowDefinition {
         .depends_on(&["register_workers"]),
     )
 }
+
+/// Helper to create initial workflow data for external worker registration
+pub fn create_external_worker_workflow_data(
+    config: WorkerConfigRequest,
+    app_context: Arc<AppContext>,
+) -> AnyWorkflowData {
+    AnyWorkflowData::ExternalWorker(ExternalWorkerWorkflowData {
+        config,
+        model_cards: Vec::new(),
+        workers: None,
+        labels: std::collections::HashMap::new(),
+        app_context: Some(app_context),
+        actual_workers: None,
+    })
+}

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

@@ -5,12 +5,12 @@ 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::{
         circuit_breaker::CircuitBreakerConfig,
         model_card::ModelCard,
+        steps::workflow_data::{AnyWorkflowData, LocalWorkerWorkflowData},
         worker::{HealthConfig, RuntimeType, WorkerType},
         BasicWorkerBuilder, ConnectionMode, DPAwareWorkerBuilder, Worker, UNKNOWN_MODEL_ID,
     },
@@ -29,13 +29,22 @@ use crate::{
 pub struct CreateLocalWorkerStep;
 
 #[async_trait]
-impl StepExecutor for CreateLocalWorkerStep {
+impl StepExecutor<AnyWorkflowData> for CreateLocalWorkerStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
+        &self,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
+    ) -> WorkflowResult<StepResult> {
-        let discovered_labels: Arc<HashMap<String, String>> =
+        let data = context.data.as_local_worker()?;
-            context.get_or_err("discovered_labels")?;
+        let config = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let connection_mode = data
+            .connection_mode
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("connection_mode".to_string()))?;
+        let discovered_labels = &data.discovered_labels;
 
         // Check if worker already exists
         if app_context
@@ -59,7 +68,7 @@ impl StepExecutor for CreateLocalWorkerStep {
         }
 
         // Merge: discovered labels first, then config labels (config takes precedence)
-        let mut final_labels = discovered_labels.as_ref().clone();
+        let mut final_labels = discovered_labels.clone();
         for (key, value) in &config_labels {
             final_labels.insert(key.clone(), value.clone());
         }
@@ -77,7 +86,7 @@ impl StepExecutor for CreateLocalWorkerStep {
         }
 
         // Create ModelCard
-        let model_card = build_model_card(&model_id, &config, &final_labels);
+        let model_card = build_model_card(&model_id, config, &final_labels);
 
         debug!(
             "Creating worker {} with {} discovered + {} config = {} final labels",
@@ -88,41 +97,39 @@ impl StepExecutor for CreateLocalWorkerStep {
         );
 
         // Parse worker type
-        let worker_type = parse_worker_type(&config);
+        let worker_type = parse_worker_type(config);
 
         // Get runtime type (for gRPC workers)
-        let runtime_type = determine_runtime_type(&connection_mode, context, &config);
+        let runtime_type = determine_runtime_type(connection_mode, data, config);
 
         // Build circuit breaker config
-        let circuit_breaker_config = build_circuit_breaker_config(&app_context);
+        let circuit_breaker_config = build_circuit_breaker_config(app_context);
 
         // Build health config
-        let health_config = build_health_config(&app_context);
+        let health_config = build_health_config(app_context);
 
         // Normalize URL
-        let normalized_url = normalize_url(&config.url, &connection_mode);
+        let normalized_url = normalize_url(&config.url, connection_mode);
 
         if normalized_url != config.url {
             debug!(
                 "Normalized worker URL: {} -> {} ({:?})",
-                config.url,
+                config.url, normalized_url, connection_mode
-                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,
+                data,
                 &normalized_url,
                 model_card,
                 worker_type,
-                &connection_mode,
+                connection_mode,
                 runtime_type,
                 circuit_breaker_config,
                 health_config,
-                &config,
+                config,
                 &final_labels,
             )?
         } else {
@@ -130,17 +137,19 @@ impl StepExecutor for CreateLocalWorkerStep {
                 &normalized_url,
                 model_card,
                 worker_type,
-                &connection_mode,
+                connection_mode,
                 runtime_type,
                 circuit_breaker_config,
                 health_config,
-                &config,
+                config,
                 &final_labels,
             )
         };
 
-        context.set("workers", workers);
+        // Update workflow data
-        context.set("labels", final_labels);
+        let data_mut = context.data.as_local_worker_mut()?;
+        data_mut.actual_workers = Some(workers);
+        data_mut.final_labels = final_labels;
         Ok(StepResult::Success)
     }
 
@@ -230,14 +239,14 @@ fn parse_worker_type(config: &WorkerConfigRequest) -> WorkerType {
 
 fn determine_runtime_type(
     connection_mode: &ConnectionMode,
-    context: &WorkflowContext,
+    data: &LocalWorkerWorkflowData,
     config: &WorkerConfigRequest,
 ) -> RuntimeType {
     if !matches!(connection_mode, ConnectionMode::Grpc { .. }) {
         return RuntimeType::Sglang;
     }
 
-    if let Some(detected_runtime) = context.get::<String>("detected_runtime_type") {
+    if let Some(ref detected_runtime) = data.detected_runtime_type {
         match detected_runtime.as_str() {
             "vllm" => RuntimeType::Vllm,
             _ => RuntimeType::Sglang,
@@ -286,7 +295,7 @@ fn normalize_url(url: &str, connection_mode: &ConnectionMode) -> String {
 
 #[allow(clippy::too_many_arguments)]
 fn create_dp_aware_workers(
-    context: &WorkflowContext,
+    data: &LocalWorkerWorkflowData,
     normalized_url: &str,
     model_card: ModelCard,
     worker_type: WorkerType,
@@ -297,7 +306,10 @@ fn create_dp_aware_workers(
     config: &WorkerConfigRequest,
     final_labels: &HashMap<String, String>,
 ) -> Result<Vec<Arc<dyn Worker>>, WorkflowError> {
-    let dp_info: Arc<DpInfo> = context.get_or_err("dp_info")?;
+    let dp_info = data
+        .dp_info
+        .as_ref()
+        .ok_or_else(|| WorkflowError::ContextValueNotFound("dp_info".to_string()))?;
 
     debug!(
         "Creating {} DP-aware workers for {} (dp_size: {})",

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

@@ -1,6 +1,6 @@
 //! Connection mode detection step.
 
-use std::{sync::Arc, time::Duration};
+use std::time::Duration;
 
 use async_trait::async_trait;
 use reqwest::Client;
@@ -8,9 +8,7 @@ use tracing::debug;
 
 use super::strip_protocol;
 use crate::{
-    app_context::AppContext,
+    core::{steps::workflow_data::AnyWorkflowData, ConnectionMode},
-    core::ConnectionMode,
-    protocols::worker_spec::WorkerConfigRequest,
     routers::grpc::client::GrpcClient,
     workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
@@ -88,10 +86,17 @@ async fn try_grpc_health_check(
 pub struct DetectConnectionModeStep;
 
 #[async_trait]
-impl StepExecutor for DetectConnectionModeStep {
+impl StepExecutor<AnyWorkflowData> for DetectConnectionModeStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
+        &self,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_local_worker()?;
+        let config = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
 
         debug!(
             "Detecting connection mode for {} (timeout: {}s, max_attempts: {})",
@@ -129,7 +134,7 @@ impl StepExecutor for DetectConnectionModeStep {
             }
         };
 
-        context.set("connection_mode", connection_mode);
+        context.data.as_local_worker_mut()?.connection_mode = Some(connection_mode);
         Ok(StepResult::Success)
     }
 

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

@@ -1,19 +1,16 @@
 //! 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::{
-    core::UNKNOWN_MODEL_ID,
+    core::{steps::workflow_data::AnyWorkflowData, UNKNOWN_MODEL_ID},
-    protocols::worker_spec::WorkerConfigRequest,
     workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 
 /// DP (Data Parallel) information for a worker.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct DpInfo {
     pub dp_size: usize,
     pub model_id: String,
@@ -44,9 +41,13 @@ pub async fn get_dp_info(url: &str, api_key: Option<&str>) -> Result<DpInfo, Str
 pub struct DiscoverDPInfoStep;
 
 #[async_trait]
-impl StepExecutor for DiscoverDPInfoStep {
+impl StepExecutor<AnyWorkflowData> for DiscoverDPInfoStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
+        &self,
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_local_worker()?;
+        let config = &data.config;
 
         if !config.dp_aware {
             debug!(
@@ -70,7 +71,7 @@ impl StepExecutor for DiscoverDPInfoStep {
             dp_info.dp_size, config.url, dp_info.model_id
         );
 
-        context.set("dp_info", dp_info);
+        context.data.as_local_worker_mut()?.dp_info = Some(dp_info);
         Ok(StepResult::Success)
     }
 

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

@@ -1,6 +1,6 @@
 //! Metadata discovery step for local workers.
 
-use std::{collections::HashMap, sync::Arc, time::Duration};
+use std::{collections::HashMap, time::Duration};
 
 use async_trait::async_trait;
 use once_cell::sync::Lazy;
@@ -11,8 +11,7 @@ use tracing::{debug, warn};
 
 use super::strip_protocol;
 use crate::{
-    core::ConnectionMode,
+    core::{steps::workflow_data::AnyWorkflowData, ConnectionMode},
-    protocols::worker_spec::WorkerConfigRequest,
     routers::grpc::client::GrpcClient,
     workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
@@ -219,17 +218,24 @@ async fn fetch_grpc_metadata(
 pub struct DiscoverMetadataStep;
 
 #[async_trait]
-impl StepExecutor for DiscoverMetadataStep {
+impl StepExecutor<AnyWorkflowData> for DiscoverMetadataStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
+        &self,
-        let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_local_worker()?;
+        let config = &data.config;
+        let connection_mode = data
+            .connection_mode
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("connection_mode".to_string()))?;
 
         debug!(
             "Discovering metadata for {} ({:?})",
-            config.url, *connection_mode
+            config.url, connection_mode
         );
 
-        let (discovered_labels, detected_runtime) = match connection_mode.as_ref() {
+        let (discovered_labels, detected_runtime) = match connection_mode {
             ConnectionMode::Http => {
                 let mut labels = HashMap::new();
 
@@ -287,16 +293,19 @@ impl StepExecutor for DiscoverMetadataStep {
             (HashMap::new(), None)
         });
 
+        let url = config.url.clone();
         debug!(
             "Discovered {} metadata labels for {}",
             discovered_labels.len(),
-            config.url
+            url
         );
 
-        context.set("discovered_labels", discovered_labels);
+        // Update workflow data
+        let data_mut = context.data.as_local_worker_mut()?;
+        data_mut.discovered_labels = discovered_labels;
         if let Some(runtime) = detected_runtime {
             debug!("Detected runtime type: {}", runtime);
-            context.set("detected_runtime_type", runtime);
+            data_mut.detected_runtime_type = Some(runtime);
         }
 
         Ok(StepResult::Success)

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

@@ -1,13 +1,11 @@
 //! Step to find a worker to update based on URL.
 
-use std::sync::Arc;
-
 use async_trait::async_trait;
 use tracing::debug;
 
 use super::find_workers_by_url;
 use crate::{
-    app_context::AppContext,
+    core::steps::workflow_data::AnyWorkflowData,
     workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 
@@ -15,31 +13,30 @@ use crate::{
 ///
 /// For DP-aware workers, finds all workers with matching URL prefix.
 /// For regular workers, finds the single worker with exact URL match.
-///
-/// Expects the following context values:
-/// - "worker_url": String - the URL of the worker to update
-/// - "dp_aware": bool - whether to find all DP-aware workers with matching prefix
-/// - "app_context": Arc<AppContext>
-///
-/// Sets the following context values:
-/// - "workers_to_update": Vec<Arc<dyn Worker>>
 pub struct FindWorkerToUpdateStep;
 
 #[async_trait]
-impl StepExecutor for FindWorkerToUpdateStep {
+impl StepExecutor<AnyWorkflowData> for FindWorkerToUpdateStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let worker_url: Arc<String> = context.get_or_err("worker_url")?;
+        &self,
-        let dp_aware: Arc<bool> = context.get_or_err("dp_aware")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_worker_update()?;
+        let worker_url = &data.worker_url;
+        let dp_aware = data.dp_aware;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
 
         let workers_to_update =
-            find_workers_by_url(&app_context.worker_registry, &worker_url, *dp_aware);
+            find_workers_by_url(&app_context.worker_registry, worker_url, dp_aware);
 
         if workers_to_update.is_empty() {
-            let error_msg = if *dp_aware {
+            let error_msg = if dp_aware {
-                format!("No workers found with prefix {}@", *worker_url)
+                format!("No workers found with prefix {}@", worker_url)
             } else {
-                format!("Worker {} not found", *worker_url)
+                format!("Worker {} not found", worker_url)
             };
             return Err(WorkflowError::StepFailed {
                 step_id: StepId::new("find_worker_to_update"),
@@ -50,10 +47,10 @@ impl StepExecutor for FindWorkerToUpdateStep {
         debug!(
             "Found {} worker(s) to update for {}",
             workers_to_update.len(),
-            *worker_url
+            worker_url
         );
 
-        context.set("workers_to_update", workers_to_update);
+        context.data.as_worker_update_mut()?.workers_to_update = Some(workers_to_update);
 
         Ok(StepResult::Success)
     }

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

@@ -1,18 +1,18 @@
 //! Step to find workers to remove based on URL.
 
-use std::{collections::HashSet, sync::Arc};
+use std::collections::HashSet;
 
 use async_trait::async_trait;
 use tracing::debug;
 
 use super::find_workers_by_url;
 use crate::{
-    app_context::AppContext,
+    core::steps::workflow_data::{AnyWorkflowData, WorkerList},
     workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 
 /// Request structure for worker removal.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub struct WorkerRemovalRequest {
     pub url: String,
     pub dp_aware: bool,
@@ -25,10 +25,17 @@ pub struct WorkerRemovalRequest {
 pub struct FindWorkersToRemoveStep;
 
 #[async_trait]
-impl StepExecutor for FindWorkersToRemoveStep {
+impl StepExecutor<AnyWorkflowData> for FindWorkersToRemoveStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let request: Arc<WorkerRemovalRequest> = context.get_or_err("removal_request")?;
+        &self,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_worker_removal()?;
+        let request = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
 
         let workers_to_remove =
             find_workers_by_url(&app_context.worker_registry, &request.url, request.dp_aware);
@@ -62,9 +69,12 @@ impl StepExecutor for FindWorkersToRemoveStep {
             .map(|w| w.model_id().to_string())
             .collect();
 
-        context.set("workers_to_remove", workers_to_remove);
+        // Update workflow data
-        context.set("worker_urls", worker_urls);
+        let data_mut = context.data.as_worker_removal_mut()?;
-        context.set("affected_models", affected_models);
+        data_mut.workers_to_remove = Some(WorkerList::from_workers(&workers_to_remove));
+        data_mut.actual_workers_to_remove = Some(workers_to_remove);
+        data_mut.worker_urls = worker_urls;
+        data_mut.affected_models = affected_models;
 
         Ok(StepResult::Success)
     }

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

@@ -36,8 +36,16 @@ pub use update_worker_properties::UpdateWorkerPropertiesStep;
 
 use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep};
 use crate::{
+    app_context::AppContext,
     config::RouterConfig,
-    core::{Worker, WorkerRegistry},
+    core::{
+        steps::workflow_data::{
+            AnyWorkflowData, LocalWorkerWorkflowData, WorkerRemovalWorkflowData,
+            WorkerUpdateWorkflowData,
+        },
+        Worker, WorkerRegistry,
+    },
+    protocols::worker_spec::{WorkerConfigRequest, WorkerUpdateRequest},
     workflow::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition},
 };
 
@@ -66,7 +74,9 @@ pub(crate) fn find_workers_by_url(
     }
 }
 
-pub fn create_local_worker_workflow(router_config: &RouterConfig) -> WorkflowDefinition {
+pub fn create_local_worker_workflow(
+    router_config: &RouterConfig,
+) -> WorkflowDefinition<AnyWorkflowData> {
     let detect_timeout = Duration::from_secs(router_config.worker_startup_timeout_secs);
 
     // Calculate max_attempts based on timeout
@@ -198,7 +208,7 @@ pub fn create_local_worker_workflow(router_config: &RouterConfig) -> WorkflowDef
 ///              │
 ///     update_remaining_policies
 /// ```
-pub fn create_worker_removal_workflow() -> WorkflowDefinition {
+pub fn create_worker_removal_workflow() -> WorkflowDefinition<AnyWorkflowData> {
     WorkflowDefinition::new("worker_removal", "Remove worker from router")
         .add_step(
             StepDefinition::new(
@@ -263,7 +273,7 @@ pub fn create_worker_removal_workflow() -> WorkflowDefinition {
 ///              │
 ///     update_policies_for_worker
 /// ```
-pub fn create_worker_update_workflow() -> WorkflowDefinition {
+pub fn create_worker_update_workflow() -> WorkflowDefinition<AnyWorkflowData> {
     WorkflowDefinition::new("worker_update", "Update worker properties")
         .add_step(
             StepDefinition::new(
@@ -304,3 +314,55 @@ pub fn create_worker_update_workflow() -> WorkflowDefinition {
             .depends_on(&["update_worker_properties"]),
         )
 }
+
+/// Helper to create initial workflow data for local worker registration
+pub fn create_local_worker_workflow_data(
+    config: WorkerConfigRequest,
+    app_context: Arc<AppContext>,
+) -> AnyWorkflowData {
+    AnyWorkflowData::LocalWorker(LocalWorkerWorkflowData {
+        config,
+        connection_mode: None,
+        discovered_labels: std::collections::HashMap::new(),
+        dp_info: None,
+        workers: None,
+        final_labels: std::collections::HashMap::new(),
+        detected_runtime_type: None,
+        app_context: Some(app_context),
+        actual_workers: None,
+    })
+}
+
+/// Helper to create initial workflow data for worker removal
+pub fn create_worker_removal_workflow_data(
+    url: String,
+    dp_aware: bool,
+    app_context: Arc<AppContext>,
+) -> AnyWorkflowData {
+    AnyWorkflowData::WorkerRemoval(WorkerRemovalWorkflowData {
+        config: WorkerRemovalRequest { url, dp_aware },
+        workers_to_remove: None,
+        worker_urls: Vec::new(),
+        affected_models: std::collections::HashSet::new(),
+        app_context: Some(app_context),
+        actual_workers_to_remove: None,
+    })
+}
+
+/// Helper to create initial workflow data for worker update
+pub fn create_worker_update_workflow_data(
+    worker_url: String,
+    update_config: WorkerUpdateRequest,
+    app_context: Arc<AppContext>,
+) -> AnyWorkflowData {
+    // Determine if this is a DP-aware update based on URL pattern
+    let dp_aware = worker_url.contains('@');
+    AnyWorkflowData::WorkerUpdate(WorkerUpdateWorkflowData {
+        config: update_config,
+        worker_url,
+        dp_aware,
+        app_context: Some(app_context),
+        workers_to_update: None,
+        updated_workers: None,
+    })
+}

sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs

@@ -1,13 +1,10 @@
 //! Tokenizer registration step for local workers.
 
-use std::{collections::HashMap, sync::Arc};
-
 use async_trait::async_trait;
 use tracing::{debug, warn};
 
 use crate::{
-    app_context::AppContext,
+    core::steps::workflow_data::AnyWorkflowData,
-    core::Worker,
     tokenizer::{factory, TokenizerRegistry},
     workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
@@ -16,11 +13,21 @@ use crate::{
 pub struct RegisterTokenizerStep;
 
 #[async_trait]
-impl StepExecutor for RegisterTokenizerStep {
+impl StepExecutor<AnyWorkflowData> for RegisterTokenizerStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
+        &self,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_local_worker()?;
+        let labels = &data.final_labels;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let workers = data
+            .actual_workers
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
 
         for worker in workers.iter() {
             let model_id = worker.model_id().to_string();
@@ -46,10 +53,11 @@ impl StepExecutor for RegisterTokenizerStep {
             let source = tokenizer_path.clone();
 
             // Load tokenizer with thread safe lock
+            let tokenizer_path_owned = tokenizer_path.clone();
             if let Err(e) = app_context
                 .tokenizer_registry
                 .load(&tokenizer_id, &model_id, &source, || async move {
-                    factory::create_tokenizer_async(&tokenizer_path.to_string())
+                    factory::create_tokenizer_async(&tokenizer_path_owned)
                         .await
                         .map_err(|e| e.to_string())
                 })

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

@@ -1,13 +1,10 @@
 //! Step to remove workers from policy registry.
 
-use std::sync::Arc;
-
 use async_trait::async_trait;
 use tracing::debug;
 
 use crate::{
-    app_context::AppContext,
+    core::steps::workflow_data::AnyWorkflowData,
-    core::Worker,
     workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 
@@ -18,11 +15,20 @@ use crate::{
 pub struct RemoveFromPolicyRegistryStep;
 
 #[async_trait]
-impl StepExecutor for RemoveFromPolicyRegistryStep {
+impl StepExecutor<AnyWorkflowData> for RemoveFromPolicyRegistryStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        &self,
-        let workers_to_remove: Arc<Vec<Arc<dyn Worker>>> =
+        context: &mut WorkflowContext<AnyWorkflowData>,
-            context.get_or_err("workers_to_remove")?;
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_worker_removal()?;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let workers_to_remove = data
+            .actual_workers_to_remove
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers_to_remove".to_string()))?;
 
         debug!(
             "Removing {} worker(s) from policy registry",

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

@@ -1,12 +1,12 @@
 //! Step to remove workers from worker registry.
 
-use std::{collections::HashSet, sync::Arc};
+use std::collections::HashSet;
 
 use async_trait::async_trait;
 use tracing::{debug, warn};
 
 use crate::{
-    app_context::AppContext,
+    core::steps::workflow_data::AnyWorkflowData,
     observability::metrics::Metrics,
     workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
@@ -17,10 +17,17 @@ use crate::{
 pub struct RemoveFromWorkerRegistryStep;
 
 #[async_trait]
-impl StepExecutor for RemoveFromWorkerRegistryStep {
+impl StepExecutor<AnyWorkflowData> for RemoveFromWorkerRegistryStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        &self,
-        let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_worker_removal()?;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let worker_urls = &data.worker_urls;
 
         debug!(
             "Removing {} worker(s) from worker registry",

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

@@ -1,13 +1,12 @@
 //! Step to update policies for updated workers.
 
-use std::{collections::HashSet, sync::Arc};
+use std::collections::HashSet;
 
 use async_trait::async_trait;
 use tracing::debug;
 
 use crate::{
-    app_context::AppContext,
+    core::steps::workflow_data::AnyWorkflowData,
-    core::Worker,
     workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 
@@ -18,10 +17,20 @@ use crate::{
 pub struct UpdatePoliciesForWorkerStep;
 
 #[async_trait]
-impl StepExecutor for UpdatePoliciesForWorkerStep {
+impl StepExecutor<AnyWorkflowData> for UpdatePoliciesForWorkerStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        &self,
-        let updated_workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("updated_workers")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_worker_update()?;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let updated_workers = data
+            .updated_workers
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("updated_workers".to_string()))?;
 
         // Collect affected models
         let affected_models: HashSet<String> = updated_workers

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

@@ -1,12 +1,10 @@
 //! Step to update cache-aware policies for remaining workers after removal.
 
-use std::{collections::HashSet, sync::Arc};
-
 use async_trait::async_trait;
 use tracing::{debug, info};
 
 use crate::{
-    app_context::AppContext,
+    core::steps::workflow_data::AnyWorkflowData,
     workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 
@@ -17,11 +15,18 @@ use crate::{
 pub struct UpdateRemainingPoliciesStep;
 
 #[async_trait]
-impl StepExecutor for UpdateRemainingPoliciesStep {
+impl StepExecutor<AnyWorkflowData> for UpdateRemainingPoliciesStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        &self,
-        let affected_models: Arc<HashSet<String>> = context.get_or_err("affected_models")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
+    ) -> WorkflowResult<StepResult> {
+        let data = context.data.as_worker_removal()?;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
+        let affected_models = &data.affected_models;
+        let worker_urls = &data.worker_urls;
 
         debug!(
             "Updating cache-aware policies for {} affected model(s)",

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

@@ -6,9 +6,7 @@ use async_trait::async_trait;
 use tracing::{debug, info};
 
 use crate::{
-    app_context::AppContext,
+    core::{steps::workflow_data::AnyWorkflowData, BasicWorkerBuilder, HealthConfig, Worker},
-    core::{BasicWorkerBuilder, HealthConfig, Worker},
-    protocols::worker_spec::WorkerUpdateRequest,
     workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 
@@ -16,23 +14,26 @@ use crate::{
 ///
 /// This step creates new worker instances with updated properties and
 /// re-registers them to replace the old workers in the registry.
-///
-/// Expects the following context values:
-/// - "update_request": WorkerUpdateRequest (from protocols::worker_spec)
-/// - "app_context": Arc<AppContext>
-/// - "workers_to_update": Vec<Arc<dyn Worker>>
-///
-/// Sets the following context values:
-/// - "updated_workers": Vec<Arc<dyn Worker>>
 pub struct UpdateWorkerPropertiesStep;
 
 #[async_trait]
-impl StepExecutor for UpdateWorkerPropertiesStep {
+impl StepExecutor<AnyWorkflowData> for UpdateWorkerPropertiesStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let request: Arc<WorkerUpdateRequest> = context.get_or_err("update_request")?;
+        &self,
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let workers_to_update: Arc<Vec<Arc<dyn Worker>>> =
+    ) -> WorkflowResult<StepResult> {
-            context.get_or_err("workers_to_update")?;
+        let data = context.data.as_worker_update()?;
+        let request = &data.config;
+        let app_context = data
+            .app_context
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
+            .clone();
+        let workers_to_update = data
+            .workers_to_update
+            .as_ref()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers_to_update".to_string()))?
+            .clone();
 
         debug!(
             "Updating properties for {} worker(s)",
@@ -136,7 +137,7 @@ impl StepExecutor for UpdateWorkerPropertiesStep {
         }
 
         // Store updated workers for subsequent steps
-        context.set("updated_workers", updated_workers);
+        context.data.as_worker_update_mut()?.updated_workers = Some(updated_workers);
 
         Ok(StepResult::Success)
     }

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

@@ -3,15 +3,16 @@ pub mod local;
 pub mod shared;
 
 pub use external::{
-    create_external_worker_workflow as create_external_worker_registration_workflow,
+    create_external_worker_workflow, create_external_worker_workflow_data, group_models_into_cards,
-    create_external_worker_workflow, group_models_into_cards, infer_model_type_from_id,
+    infer_model_type_from_id, CreateExternalWorkersStep, DiscoverModelsStep, ModelInfo,
-    CreateExternalWorkersStep, DiscoverModelsStep, ModelInfo, ModelsResponse,
+    ModelsResponse,
 };
 pub use local::{
-    create_local_worker_workflow as create_worker_registration_workflow,
+    create_local_worker_workflow, create_local_worker_workflow_data,
-    create_local_worker_workflow, create_worker_removal_workflow, create_worker_update_workflow,
+    create_worker_removal_workflow, create_worker_removal_workflow_data,
-    CreateLocalWorkerStep, DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep,
+    create_worker_update_workflow, create_worker_update_workflow_data, CreateLocalWorkerStep,
-    DpInfo, FindWorkerToUpdateStep, FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep,
+    DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep, DpInfo,
+    FindWorkerToUpdateStep, FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep,
     RemoveFromWorkerRegistryStep, UpdatePoliciesForWorkerStep, UpdateRemainingPoliciesStep,
     UpdateWorkerPropertiesStep, WorkerRemovalRequest,
 };

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

@@ -1,13 +1,11 @@
 //! Unified worker activation step.
 
-use std::sync::Arc;
-
 use async_trait::async_trait;
 use tracing::info;
 
 use crate::{
-    core::Worker,
+    core::steps::workflow_data::AnyWorkflowData,
-    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult},
+    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 
 /// Unified step to activate workers by marking them as healthy.
@@ -16,9 +14,15 @@ use crate::{
 pub struct ActivateWorkersStep;
 
 #[async_trait]
-impl StepExecutor for ActivateWorkersStep {
+impl StepExecutor<AnyWorkflowData> for ActivateWorkersStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
+        &self,
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let workers = context
+            .data
+            .get_actual_workers()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
 
         for worker in workers.iter() {
             worker.set_healthy(true);
@@ -29,7 +33,7 @@ impl StepExecutor for ActivateWorkersStep {
         Ok(StepResult::Success)
     }
 
-    fn is_retryable(&self, _error: &crate::workflow::WorkflowError) -> bool {
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
         false
     }
 }

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

@@ -6,10 +6,9 @@ use async_trait::async_trait;
 use tracing::debug;
 
 use crate::{
-    app_context::AppContext,
+    core::steps::workflow_data::AnyWorkflowData,
-    core::Worker,
     observability::metrics::Metrics,
-    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult},
+    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
 };
 
 /// Unified step to register workers in the registry.
@@ -19,10 +18,21 @@ use crate::{
 pub struct RegisterWorkersStep;
 
 #[async_trait]
-impl StepExecutor for RegisterWorkersStep {
+impl StepExecutor<AnyWorkflowData> for RegisterWorkersStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        &self,
-        let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
+        let app_context = context
+            .data
+            .get_app_context()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
+            .clone();
+
+        let workers = context
+            .data
+            .get_actual_workers()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
 
         let mut worker_ids = Vec::with_capacity(workers.len());
 
@@ -75,11 +85,18 @@ impl StepExecutor for RegisterWorkersStep {
             );
         }
 
-        context.set("worker_ids", worker_ids);
+        // Note: worker_ids are stored for potential future use but not persisted
+        // as they are internal registry identifiers
+        debug!(
+            "Registered {} workers with IDs: {:?}",
+            worker_ids.len(),
+            worker_ids
+        );
+
         Ok(StepResult::Success)
     }
 
-    fn is_retryable(&self, _error: &crate::workflow::WorkflowError) -> bool {
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
         false
     }
 }

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

@@ -1,14 +1,13 @@
 //! Unified policy update step.
 
-use std::{collections::HashMap, sync::Arc};
+use std::sync::Arc;
 
 use async_trait::async_trait;
 use tracing::{debug, warn};
 
 use crate::{
-    app_context::AppContext,
+    core::{steps::workflow_data::AnyWorkflowData, Worker},
-    core::Worker,
+    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
-    workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult},
 };
 
 /// Unified step to update policy registry for registered workers.
@@ -82,11 +81,26 @@ impl UpdatePoliciesStep {
 }
 
 #[async_trait]
-impl StepExecutor for UpdatePoliciesStep {
+impl StepExecutor<AnyWorkflowData> for UpdatePoliciesStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
-        let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
+        &self,
-        let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
+        context: &mut WorkflowContext<AnyWorkflowData>,
-        let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
+    ) -> WorkflowResult<StepResult> {
+        let app_context = context
+            .data
+            .get_app_context()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
+            .clone();
+
+        let workers = context
+            .data
+            .get_actual_workers()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
+
+        let labels = context
+            .data
+            .get_labels()
+            .ok_or_else(|| WorkflowError::ContextValueNotFound("labels".to_string()))?;
 
         let policy_hint = labels.get("policy").map(|s| s.as_str());
 
@@ -139,7 +153,7 @@ impl StepExecutor for UpdatePoliciesStep {
         Ok(StepResult::Success)
     }
 
-    fn is_retryable(&self, _error: &crate::workflow::WorkflowError) -> bool {
+    fn is_retryable(&self, _error: &WorkflowError) -> bool {
         false
     }
 }

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

@@ -0,0 +1,559 @@
+//! Typed workflow data structures
+//!
+//! This module defines the typed data structures for all workflows, enabling
+//! compile-time type safety and state persistence.
+
+use std::{collections::HashMap, sync::Arc};
+
+use serde::{Deserialize, Serialize};
+
+use super::{
+    mcp_registration::McpServerConfigRequest, tokenizer_registration::TokenizerConfigRequest,
+    wasm_module_registration::WasmModuleConfigRequest,
+    wasm_module_removal::WasmModuleRemovalRequest, worker::local::WorkerRemovalRequest,
+};
+/// Re-export the protocol types for convenience
+pub use crate::protocols::worker_spec::{
+    WorkerConfigRequest, WorkerUpdateRequest as ProtocolUpdateRequest,
+};
+use crate::{
+    app_context::AppContext,
+    core::{model_card::ModelCard, Worker},
+    protocols::worker_spec::{
+        WorkerConfigRequest as ProtocolWorkerConfigRequest,
+        WorkerUpdateRequest as ProtocolWorkerUpdateRequest,
+    },
+    workflow::{WorkflowData, WorkflowError},
+};
+
+/// Wrapper for worker list that can be serialized
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct WorkerList {
+    /// Worker URLs (we can't serialize Arc<dyn Worker>, so we store URLs)
+    pub worker_urls: Vec<String>,
+}
+
+impl WorkerList {
+    pub fn new() -> Self {
+        Self {
+            worker_urls: Vec::new(),
+        }
+    }
+
+    pub fn from_workers(workers: &[Arc<dyn Worker>]) -> Self {
+        Self {
+            worker_urls: workers.iter().map(|w| w.url().to_string()).collect(),
+        }
+    }
+}
+
+// ============================================================================
+// Workflow-specific data types
+// ============================================================================
+
+/// Data for tokenizer registration workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TokenizerWorkflowData {
+    pub config: TokenizerConfigRequest,
+    pub vocab_size: Option<usize>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+}
+
+impl WorkflowData for TokenizerWorkflowData {
+    fn workflow_type() -> &'static str {
+        "tokenizer_registration"
+    }
+}
+
+impl TokenizerWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    ///
+    /// Call this after deserializing workflow state to ensure runtime fields
+    /// have been repopulated.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for local worker registration workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LocalWorkerWorkflowData {
+    pub config: ProtocolWorkerConfigRequest,
+    pub connection_mode: Option<crate::core::ConnectionMode>,
+    pub discovered_labels: HashMap<String, String>,
+    pub dp_info: Option<super::worker::local::DpInfo>,
+    pub workers: Option<WorkerList>,
+    pub final_labels: HashMap<String, String>,
+    /// Detected runtime type (for gRPC workers)
+    pub detected_runtime_type: Option<String>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+    /// Actual worker objects (transient, not serialized)
+    #[serde(skip, default)]
+    pub actual_workers: Option<Vec<Arc<dyn Worker>>>,
+}
+
+impl WorkflowData for LocalWorkerWorkflowData {
+    fn workflow_type() -> &'static str {
+        "local_worker_registration"
+    }
+}
+
+impl LocalWorkerWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for external worker registration workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ExternalWorkerWorkflowData {
+    pub config: ProtocolWorkerConfigRequest,
+    /// Discovered model cards from /v1/models endpoint
+    pub model_cards: Vec<ModelCard>,
+    pub workers: Option<WorkerList>,
+    /// Labels for policies (derived from config)
+    pub labels: HashMap<String, String>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+    /// Actual worker objects (transient, not serialized)
+    #[serde(skip, default)]
+    pub actual_workers: Option<Vec<Arc<dyn Worker>>>,
+}
+
+impl WorkflowData for ExternalWorkerWorkflowData {
+    fn workflow_type() -> &'static str {
+        "external_worker_registration"
+    }
+}
+
+impl ExternalWorkerWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for worker removal workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WorkerRemovalWorkflowData {
+    pub config: WorkerRemovalRequest,
+    pub workers_to_remove: Option<WorkerList>,
+    /// URLs of workers being removed
+    pub worker_urls: Vec<String>,
+    /// Model IDs affected by the removal
+    pub affected_models: std::collections::HashSet<String>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+    /// Actual worker objects to remove (transient, not serialized)
+    #[serde(skip, default)]
+    pub actual_workers_to_remove: Option<Vec<Arc<dyn Worker>>>,
+}
+
+impl WorkflowData for WorkerRemovalWorkflowData {
+    fn workflow_type() -> &'static str {
+        "worker_removal"
+    }
+}
+
+impl WorkerRemovalWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for worker update workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WorkerUpdateWorkflowData {
+    pub config: ProtocolWorkerUpdateRequest,
+    /// URL of worker(s) to update
+    pub worker_url: String,
+    /// Whether to update all DP-aware workers with matching prefix
+    pub dp_aware: bool,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+    /// Workers to update (transient, not serialized)
+    #[serde(skip, default)]
+    pub workers_to_update: Option<Vec<Arc<dyn Worker>>>,
+    /// Updated worker objects (transient, not serialized)
+    #[serde(skip, default)]
+    pub updated_workers: Option<Vec<Arc<dyn Worker>>>,
+}
+
+impl WorkflowData for WorkerUpdateWorkflowData {
+    fn workflow_type() -> &'static str {
+        "worker_update"
+    }
+}
+
+impl WorkerUpdateWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for MCP server registration workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct McpWorkflowData {
+    pub config: McpServerConfigRequest,
+    pub validated: bool,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+    /// Connected MCP client (transient, not serialized)
+    #[serde(skip, default)]
+    pub mcp_client: Option<Arc<rmcp::service::RunningService<rmcp::RoleClient, ()>>>,
+}
+
+impl WorkflowData for McpWorkflowData {
+    fn workflow_type() -> &'static str {
+        "mcp_registration"
+    }
+}
+
+impl McpWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for WASM module registration workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WasmRegistrationWorkflowData {
+    pub config: WasmModuleConfigRequest,
+    pub wasm_bytes: Option<Vec<u8>>,
+    /// SHA256 hash of the module file (32 bytes)
+    pub sha256_hash: Option<[u8; 32]>,
+    /// File size in bytes
+    pub file_size_bytes: Option<u64>,
+    /// UUID assigned to the registered module
+    pub module_uuid: Option<uuid::Uuid>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+}
+
+impl WorkflowData for WasmRegistrationWorkflowData {
+    fn workflow_type() -> &'static str {
+        "wasm_module_registration"
+    }
+}
+
+impl WasmRegistrationWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Data for WASM module removal workflow
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WasmRemovalWorkflowData {
+    pub config: WasmModuleRemovalRequest,
+    pub module_id: Option<String>,
+    /// Application context (transient, must be re-initialized after deserialization)
+    #[serde(skip, default)]
+    pub app_context: Option<Arc<AppContext>>,
+}
+
+impl WorkflowData for WasmRemovalWorkflowData {
+    fn workflow_type() -> &'static str {
+        "wasm_module_removal"
+    }
+}
+
+impl WasmRemovalWorkflowData {
+    /// Validate that all transient fields are properly initialized.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        if self.app_context.is_none() {
+            return Err(WorkflowError::ContextValueNotFound(
+                "app_context not initialized after deserialization".into(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+// ============================================================================
+// Unified enum for all workflow types
+// ============================================================================
+
+/// Macro to generate type-safe accessor methods for AnyWorkflowData variants.
+///
+/// This reduces boilerplate and ensures consistent error handling across all accessors.
+macro_rules! impl_workflow_accessor {
+    ($fn_name:ident, $fn_name_mut:ident, $variant:ident, $ty:ty, $type_name:expr) => {
+        /// Extract the inner data, returning an error if this is a different variant.
+        #[must_use = "this returns the result of the operation, without modifying the original"]
+        pub fn $fn_name(&self) -> Result<&$ty, WorkflowError> {
+            match self {
+                AnyWorkflowData::$variant(data) => Ok(data),
+                _ => Err(WorkflowError::TypeMismatch {
+                    expected: $type_name,
+                    actual: self.concrete_type(),
+                }),
+            }
+        }
+
+        /// Extract the inner data mutably, returning an error if this is a different variant.
+        pub fn $fn_name_mut(&mut self) -> Result<&mut $ty, WorkflowError> {
+            // Store the type name before the mutable borrow
+            let actual = self.concrete_type();
+            match self {
+                AnyWorkflowData::$variant(data) => Ok(data),
+                _ => Err(WorkflowError::TypeMismatch {
+                    expected: $type_name,
+                    actual,
+                }),
+            }
+        }
+    };
+}
+
+/// Macro to generate From implementations for AnyWorkflowData variants.
+macro_rules! impl_from_workflow_data {
+    ($variant:ident, $ty:ty) => {
+        impl From<$ty> for AnyWorkflowData {
+            fn from(data: $ty) -> Self {
+                AnyWorkflowData::$variant(data)
+            }
+        }
+    };
+}
+
+/// Unified workflow data enum covering all workflow types.
+///
+/// This allows a single `WorkflowEngine<AnyWorkflowData>` to handle all workflows
+/// while maintaining type safety at the step level.
+///
+/// # Type Erasure
+///
+/// `AnyWorkflowData` implements `WorkflowData` with `workflow_type()` returning `"any"`.
+/// This is intentional: the static method cannot know the runtime variant. Use
+/// [`concrete_type()`](Self::concrete_type) to get the actual workflow type at runtime.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum AnyWorkflowData {
+    Tokenizer(TokenizerWorkflowData),
+    LocalWorker(LocalWorkerWorkflowData),
+    ExternalWorker(ExternalWorkerWorkflowData),
+    WorkerRemoval(WorkerRemovalWorkflowData),
+    WorkerUpdate(WorkerUpdateWorkflowData),
+    Mcp(McpWorkflowData),
+    WasmRegistration(WasmRegistrationWorkflowData),
+    WasmRemoval(WasmRemovalWorkflowData),
+}
+
+impl WorkflowData for AnyWorkflowData {
+    /// Returns `"any"` as this is a type-erased container.
+    ///
+    /// Use [`concrete_type()`](Self::concrete_type) to get the actual workflow type at runtime.
+    fn workflow_type() -> &'static str {
+        "any"
+    }
+}
+
+// Generate From implementations for ergonomic construction
+impl_from_workflow_data!(Tokenizer, TokenizerWorkflowData);
+impl_from_workflow_data!(LocalWorker, LocalWorkerWorkflowData);
+impl_from_workflow_data!(ExternalWorker, ExternalWorkerWorkflowData);
+impl_from_workflow_data!(WorkerRemoval, WorkerRemovalWorkflowData);
+impl_from_workflow_data!(WorkerUpdate, WorkerUpdateWorkflowData);
+impl_from_workflow_data!(Mcp, McpWorkflowData);
+impl_from_workflow_data!(WasmRegistration, WasmRegistrationWorkflowData);
+impl_from_workflow_data!(WasmRemoval, WasmRemovalWorkflowData);
+
+impl AnyWorkflowData {
+    /// Get the concrete workflow type name at runtime.
+    ///
+    /// Unlike the static `workflow_type()` method, this returns the actual
+    /// type of the contained workflow data.
+    #[must_use]
+    pub fn concrete_type(&self) -> &'static str {
+        match self {
+            AnyWorkflowData::Tokenizer(_) => TokenizerWorkflowData::workflow_type(),
+            AnyWorkflowData::LocalWorker(_) => LocalWorkerWorkflowData::workflow_type(),
+            AnyWorkflowData::ExternalWorker(_) => ExternalWorkerWorkflowData::workflow_type(),
+            AnyWorkflowData::WorkerRemoval(_) => WorkerRemovalWorkflowData::workflow_type(),
+            AnyWorkflowData::WorkerUpdate(_) => WorkerUpdateWorkflowData::workflow_type(),
+            AnyWorkflowData::Mcp(_) => McpWorkflowData::workflow_type(),
+            AnyWorkflowData::WasmRegistration(_) => WasmRegistrationWorkflowData::workflow_type(),
+            AnyWorkflowData::WasmRemoval(_) => WasmRemovalWorkflowData::workflow_type(),
+        }
+    }
+
+    // Generate all accessor methods using the macro
+    impl_workflow_accessor!(
+        as_tokenizer,
+        as_tokenizer_mut,
+        Tokenizer,
+        TokenizerWorkflowData,
+        "tokenizer_registration"
+    );
+    impl_workflow_accessor!(
+        as_local_worker,
+        as_local_worker_mut,
+        LocalWorker,
+        LocalWorkerWorkflowData,
+        "local_worker_registration"
+    );
+    impl_workflow_accessor!(
+        as_external_worker,
+        as_external_worker_mut,
+        ExternalWorker,
+        ExternalWorkerWorkflowData,
+        "external_worker_registration"
+    );
+    impl_workflow_accessor!(
+        as_worker_removal,
+        as_worker_removal_mut,
+        WorkerRemoval,
+        WorkerRemovalWorkflowData,
+        "worker_removal"
+    );
+    impl_workflow_accessor!(
+        as_worker_update,
+        as_worker_update_mut,
+        WorkerUpdate,
+        WorkerUpdateWorkflowData,
+        "worker_update"
+    );
+    impl_workflow_accessor!(as_mcp, as_mcp_mut, Mcp, McpWorkflowData, "mcp_registration");
+    impl_workflow_accessor!(
+        as_wasm_registration,
+        as_wasm_registration_mut,
+        WasmRegistration,
+        WasmRegistrationWorkflowData,
+        "wasm_module_registration"
+    );
+    impl_workflow_accessor!(
+        as_wasm_removal,
+        as_wasm_removal_mut,
+        WasmRemoval,
+        WasmRemovalWorkflowData,
+        "wasm_module_removal"
+    );
+
+    // ========================================================================
+    // Helper methods for shared worker steps
+    // ========================================================================
+
+    /// Get app_context from any workflow data type that has it.
+    #[must_use]
+    pub fn get_app_context(&self) -> Option<&Arc<AppContext>> {
+        match self {
+            AnyWorkflowData::Tokenizer(d) => d.app_context.as_ref(),
+            AnyWorkflowData::LocalWorker(d) => d.app_context.as_ref(),
+            AnyWorkflowData::ExternalWorker(d) => d.app_context.as_ref(),
+            AnyWorkflowData::WorkerRemoval(d) => d.app_context.as_ref(),
+            AnyWorkflowData::WorkerUpdate(d) => d.app_context.as_ref(),
+            AnyWorkflowData::Mcp(d) => d.app_context.as_ref(),
+            AnyWorkflowData::WasmRegistration(d) => d.app_context.as_ref(),
+            AnyWorkflowData::WasmRemoval(d) => d.app_context.as_ref(),
+        }
+    }
+
+    /// Get actual workers from local or external worker workflows.
+    #[must_use]
+    pub fn get_actual_workers(&self) -> Option<&Vec<Arc<dyn Worker>>> {
+        match self {
+            AnyWorkflowData::LocalWorker(d) => d.actual_workers.as_ref(),
+            AnyWorkflowData::ExternalWorker(d) => d.actual_workers.as_ref(),
+            _ => None,
+        }
+    }
+
+    /// Set actual workers for local or external worker workflows.
+    pub fn set_actual_workers(
+        &mut self,
+        workers: Vec<Arc<dyn Worker>>,
+    ) -> Result<(), WorkflowError> {
+        match self {
+            AnyWorkflowData::LocalWorker(d) => {
+                d.workers = Some(WorkerList::from_workers(&workers));
+                d.actual_workers = Some(workers);
+                Ok(())
+            }
+            AnyWorkflowData::ExternalWorker(d) => {
+                d.workers = Some(WorkerList::from_workers(&workers));
+                d.actual_workers = Some(workers);
+                Ok(())
+            }
+            _ => Err(WorkflowError::TypeMismatch {
+                expected: "LocalWorker or ExternalWorker",
+                actual: self.concrete_type(),
+            }),
+        }
+    }
+
+    /// Get labels for policy configuration (from local or external worker workflows).
+    #[must_use]
+    pub fn get_labels(&self) -> Option<&HashMap<String, String>> {
+        match self {
+            AnyWorkflowData::LocalWorker(d) => Some(&d.final_labels),
+            AnyWorkflowData::ExternalWorker(d) => Some(&d.labels),
+            _ => None,
+        }
+    }
+
+    /// Validate that all transient fields are properly initialized.
+    ///
+    /// Call this after deserializing workflow state to ensure runtime fields
+    /// have been repopulated.
+    pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
+        match self {
+            AnyWorkflowData::Tokenizer(d) => d.validate_initialized(),
+            AnyWorkflowData::LocalWorker(d) => d.validate_initialized(),
+            AnyWorkflowData::ExternalWorker(d) => d.validate_initialized(),
+            AnyWorkflowData::WorkerRemoval(d) => d.validate_initialized(),
+            AnyWorkflowData::WorkerUpdate(d) => d.validate_initialized(),
+            AnyWorkflowData::Mcp(d) => d.validate_initialized(),
+            AnyWorkflowData::WasmRegistration(d) => d.validate_initialized(),
+            AnyWorkflowData::WasmRemoval(d) => d.validate_initialized(),
+        }
+    }
+}

sgl-model-gateway/src/server.rs

@@ -20,14 +20,14 @@ use tokio::{signal, spawn};
 use tracing::{debug, error, info, warn, Level};
 
 use crate::{
-    app_context::AppContext,
+    app_context::{AppContext, AppWorkflowEngine},
     config::{RouterConfig, RoutingMode},
     core::{
         job_queue::{JobQueue, JobQueueConfig},
         steps::{
-            create_external_worker_registration_workflow, create_mcp_registration_workflow,
+            create_external_worker_workflow, create_local_worker_workflow,
-            create_tokenizer_registration_workflow, create_wasm_module_registration_workflow,
+            create_mcp_registration_workflow, create_tokenizer_registration_workflow,
-            create_wasm_module_removal_workflow, create_worker_registration_workflow,
+            create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
             create_worker_removal_workflow, create_worker_update_workflow,
         },
         worker::WorkerType,
@@ -56,7 +56,7 @@ use crate::{
     routers::{conversations, parse, router_manager::RouterManager, tokenize, RouterTrait},
     service_discovery::{start_service_discovery, ServiceDiscoveryConfig},
     wasm::route::{add_wasm_module, list_wasm_modules, remove_wasm_module},
-    workflow::{LoggingSubscriber, WorkflowEngine},
+    workflow::LoggingSubscriber,
 };
 #[derive(Clone)]
 pub struct AppState {
@@ -731,7 +731,7 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
         .expect("JobQueue should only be initialized once");
 
     // Initialize workflow engine and register workflows
-    let engine = Arc::new(WorkflowEngine::new());
+    let engine = Arc::new(AppWorkflowEngine::new());
 
     engine
         .event_bus()
@@ -739,10 +739,10 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
         .await;
 
     engine
-        .register_workflow(create_worker_registration_workflow(&config.router_config))
+        .register_workflow(create_local_worker_workflow(&config.router_config))
-        .expect("worker_registration workflow should be valid");
+        .expect("local_worker_registration workflow should be valid");
     engine
-        .register_workflow(create_external_worker_registration_workflow())
+        .register_workflow(create_external_worker_workflow())
         .expect("external_worker_registration workflow should be valid");
     engine
         .register_workflow(create_worker_removal_workflow())

sgl-model-gateway/src/workflow/definition.rs

@@ -8,25 +8,25 @@ use std::{
 
 use super::{
     executor::StepExecutor,
-    types::{FailureAction, RetryPolicy, StepId, WorkflowId},
+    types::{FailureAction, RetryPolicy, StepId, WorkflowData, WorkflowId},
 };
 
 /// Definition of a single step within a workflow
-pub struct StepDefinition {
+pub struct StepDefinition<D: WorkflowData> {
     pub id: StepId,
     pub name: String,
-    pub executor: Arc<dyn StepExecutor>,
+    pub executor: Arc<dyn StepExecutor<D>>,
     pub retry_policy: Option<RetryPolicy>,
     pub timeout: Option<Duration>,
     pub on_failure: FailureAction,
     pub depends_on: Vec<StepId>,
 }
 
-impl StepDefinition {
+impl<D: WorkflowData> StepDefinition<D> {
     pub fn new(
         id: impl Into<String>,
         name: impl Into<String>,
-        executor: Arc<dyn StepExecutor>,
+        executor: Arc<dyn StepExecutor<D>>,
     ) -> Self {
         Self {
             id: StepId::new(id.into()),
@@ -64,10 +64,10 @@ impl StepDefinition {
 }
 
 /// Complete workflow definition
-pub struct WorkflowDefinition {
+pub struct WorkflowDefinition<D: WorkflowData> {
     pub id: WorkflowId,
     pub name: String,
-    pub steps: Vec<StepDefinition>,
+    pub steps: Vec<StepDefinition<D>>,
     pub default_retry_policy: RetryPolicy,
     pub default_timeout: Duration,
     /// Pre-computed reverse dependencies: step_id -> indices of steps that depend on it
@@ -76,7 +76,7 @@ pub struct WorkflowDefinition {
     initial_step_indices: Vec<usize>,
 }
 
-impl WorkflowDefinition {
+impl<D: WorkflowData> WorkflowDefinition<D> {
     pub fn new(id: impl Into<String>, name: impl Into<String>) -> Self {
         Self {
             id: WorkflowId::new(id.into()),
@@ -89,7 +89,7 @@ impl WorkflowDefinition {
         }
     }
 
-    pub fn add_step(mut self, step: StepDefinition) -> Self {
+    pub fn add_step(mut self, step: StepDefinition<D>) -> Self {
         self.steps.push(step);
         self
     }
@@ -105,14 +105,14 @@ impl WorkflowDefinition {
     }
 
     /// Get the retry policy for a step (step-specific or default)
-    pub fn get_retry_policy<'a>(&'a self, step: &'a StepDefinition) -> &'a RetryPolicy {
+    pub fn get_retry_policy<'a>(&'a self, step: &'a StepDefinition<D>) -> &'a RetryPolicy {
         step.retry_policy
             .as_ref()
             .unwrap_or(&self.default_retry_policy)
     }
 
     /// Get the timeout for a step (step-specific or default)
-    pub fn get_timeout(&self, step: &StepDefinition) -> Duration {
+    pub fn get_timeout(&self, step: &StepDefinition<D>) -> Duration {
         step.timeout.unwrap_or(self.default_timeout)
     }
 
@@ -124,7 +124,7 @@ impl WorkflowDefinition {
     /// On success, pre-computes reverse dependencies for O(1) dependent lookup.
     pub fn validate(&mut self) -> Result<(), String> {
         // Build HashMap for O(1) lookup instead of O(n) linear search
-        let steps_map: HashMap<&StepId, &StepDefinition> =
+        let steps_map: HashMap<&StepId, &StepDefinition<D>> =
             self.steps.iter().map(|s| (&s.id, s)).collect();
 
         // Check all dependencies exist
@@ -177,7 +177,7 @@ impl WorkflowDefinition {
     /// DFS helper for cycle detection with O(1) HashMap lookup
     fn has_cycle<'a>(
         step_id: &'a StepId,
-        steps_map: &HashMap<&'a StepId, &'a StepDefinition>,
+        steps_map: &HashMap<&'a StepId, &'a StepDefinition<D>>,
         visited: &mut HashSet<&'a StepId>,
         rec_stack: &mut HashSet<&'a StepId>,
     ) -> bool {

sgl-model-gateway/src/workflow/engine.rs

@@ -6,6 +6,7 @@
 
 use std::{
     collections::{HashMap, HashSet, VecDeque},
+    marker::PhantomData,
     sync::{
         atomic::{AtomicUsize, Ordering},
         Arc,
@@ -24,7 +25,7 @@ use tokio::{
 use super::{
     definition::{StepDefinition, WorkflowDefinition},
     event::{EventBus, WorkflowEvent},
-    state::WorkflowStateStore,
+    state::{InMemoryStore, StateStore},
     types::*,
 };
 
@@ -101,6 +102,11 @@ impl Backoff for LinearBackoff {
 
 /// Main workflow execution engine
 ///
+/// # Type Parameters
+///
+/// * `D` - The workflow data type that implements `WorkflowData`
+/// * `S` - The state store implementation (defaults to `InMemoryStore<D>`)
+///
 /// # Graceful Shutdown
 ///
 /// The engine supports graceful shutdown via [`shutdown()`](Self::shutdown):
@@ -115,9 +121,9 @@ impl Backoff for LinearBackoff {
 ///     engine.force_cancel_all().await;
 /// }
 /// ```
-pub struct WorkflowEngine {
+pub struct WorkflowEngine<D: WorkflowData, S: StateStore<D> = InMemoryStore<D>> {
-    definitions: Arc<RwLock<HashMap<WorkflowId, Arc<WorkflowDefinition>>>>,
+    definitions: Arc<RwLock<HashMap<WorkflowId, Arc<WorkflowDefinition<D>>>>>,
-    state_store: WorkflowStateStore,
+    state_store: S,
     event_bus: Arc<EventBus>,
     /// Shutdown signal sender - when true, engine is shutting down
     shutdown_tx: Arc<watch::Sender<bool>>,
@@ -125,18 +131,27 @@ pub struct WorkflowEngine {
     shutdown_rx: watch::Receiver<bool>,
     /// Count of active workflow executions
     active_workflows: Arc<AtomicUsize>,
+    _phantom: PhantomData<D>,
 }
 
-impl WorkflowEngine {
+impl<D: WorkflowData> WorkflowEngine<D, InMemoryStore<D>> {
     pub fn new() -> Self {
+        Self::with_store(InMemoryStore::new())
+    }
+}
+
+impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
+    /// Create a new workflow engine with a custom state store
+    pub fn with_store(state_store: S) -> Self {
         let (shutdown_tx, shutdown_rx) = watch::channel(false);
         Self {
             definitions: Arc::new(RwLock::new(HashMap::new())),
-            state_store: WorkflowStateStore::new(),
+            state_store,
             event_bus: Arc::new(EventBus::new()),
             shutdown_tx: Arc::new(shutdown_tx),
             shutdown_rx,
             active_workflows: Arc::new(AtomicUsize::new(0)),
+            _phantom: PhantomData,
         }
     }
 
@@ -284,7 +299,7 @@ impl WorkflowEngine {
     }
 
     /// Register a workflow definition
-    pub fn register_workflow(&self, mut definition: WorkflowDefinition) -> Result<(), String> {
+    pub fn register_workflow(&self, mut definition: WorkflowDefinition<D>) -> Result<(), String> {
         // Validate DAG and build dependency graph once at registration
         definition.validate()?;
 
@@ -299,7 +314,7 @@ impl WorkflowEngine {
     }
 
     /// Get the state store
-    pub fn state_store(&self) -> &WorkflowStateStore {
+    pub fn state_store(&self) -> &S {
         &self.state_store
     }
 
@@ -309,7 +324,7 @@ impl WorkflowEngine {
     pub async fn start_workflow(
         &self,
         definition_id: WorkflowId,
-        context: WorkflowContext,
+        data: D,
     ) -> WorkflowResult<WorkflowInstanceId> {
         // Guard increments counter and decrements on drop unless committed.
         // This handles all error paths automatically.
@@ -326,10 +341,9 @@ impl WorkflowEngine {
             .cloned()
             .ok_or_else(|| WorkflowError::DefinitionNotFound(definition_id.clone()))?;
 
-        let instance_id = context.instance_id;
+        let instance_id = WorkflowInstanceId::new();
-        let mut state = WorkflowState::new(instance_id, definition_id.clone());
+        let mut state = WorkflowState::new(instance_id, definition_id.clone(), data);
         state.status = WorkflowStatus::Running;
-        state.context = context;
 
         for step in &definition.steps {
             state
@@ -369,7 +383,7 @@ impl WorkflowEngine {
     async fn execute_workflow(
         &self,
         instance_id: WorkflowInstanceId,
-        definition: Arc<WorkflowDefinition>,
+        definition: Arc<WorkflowDefinition<D>>,
     ) -> WorkflowResult<()> {
         let start_time = std::time::Instant::now();
         let step_count = definition.steps.len();
@@ -572,8 +586,8 @@ impl WorkflowEngine {
     async fn execute_step_with_retry(
         &self,
         instance_id: WorkflowInstanceId,
-        step: &StepDefinition,
+        step: &StepDefinition<D>,
-        definition: &WorkflowDefinition,
+        definition: &WorkflowDefinition<D>,
     ) -> WorkflowResult<StepResult> {
         let retry_policy = definition.get_retry_policy(step);
         let step_timeout = definition.get_timeout(step);
@@ -624,7 +638,7 @@ impl WorkflowEngine {
             let step_duration = step_start.elapsed();
 
             self.state_store.update(instance_id, |s| {
-                s.context = std::mem::replace(&mut context, WorkflowContext::new(instance_id));
+                s.context = context.clone();
             })?;
 
             match result {
@@ -763,7 +777,7 @@ impl WorkflowEngine {
     }
 
     /// Get workflow status
-    pub fn get_status(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState> {
+    pub fn get_status(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>> {
         self.state_store.load(instance_id)
     }
 
@@ -776,6 +790,7 @@ impl WorkflowEngine {
             shutdown_tx: Arc::clone(&self.shutdown_tx),
             shutdown_rx: self.shutdown_rx.clone(),
             active_workflows: Arc::clone(&self.active_workflows),
+            _phantom: PhantomData,
         }
     }
 }
@@ -794,13 +809,13 @@ impl Drop for ActiveWorkflowGuard {
 
 /// RAII guard for start_workflow that increments on creation and decrements on drop
 /// unless commit() is called. Handles all error paths automatically.
-struct StartGuard<'a> {
+struct StartGuard<'a, D: WorkflowData, S: StateStore<D> + 'static> {
-    engine: &'a WorkflowEngine,
+    engine: &'a WorkflowEngine<D, S>,
     committed: bool,
 }
 
-impl<'a> StartGuard<'a> {
+impl<'a, D: WorkflowData, S: StateStore<D> + 'static> StartGuard<'a, D, S> {
-    fn new(engine: &'a WorkflowEngine) -> Self {
+    fn new(engine: &'a WorkflowEngine<D, S>) -> Self {
         engine.active_workflows.fetch_add(1, Ordering::AcqRel);
         Self {
             engine,
@@ -813,7 +828,7 @@ impl<'a> StartGuard<'a> {
     }
 }
 
-impl Drop for StartGuard<'_> {
+impl<D: WorkflowData, S: StateStore<D> + 'static> Drop for StartGuard<'_, D, S> {
     fn drop(&mut self) {
         if !self.committed {
             self.engine.workflow_finished();
@@ -821,23 +836,35 @@ impl Drop for StartGuard<'_> {
     }
 }
 
-impl Clone for WorkflowEngine {
+/// Clone implementation for internal use.
+///
+/// **Note**: This creates a shallow clone that shares state with the original engine.
+/// Both engines will share the same:
+/// - Workflow definitions
+/// - State store
+/// - Event bus
+/// - Shutdown signal
+/// - Active workflow counter
+///
+/// This is intentional for spawning async tasks that need access to the engine.
+/// For most use cases, prefer sharing the engine via `Arc<WorkflowEngine>` rather
+/// than cloning.
+impl<D: WorkflowData, S: StateStore<D> + 'static> Clone for WorkflowEngine<D, S> {
     fn clone(&self) -> Self {
         self.clone_for_execution()
     }
 }
 
-impl Default for WorkflowEngine {
+impl<D: WorkflowData> Default for WorkflowEngine<D, InMemoryStore<D>> {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl std::fmt::Debug for WorkflowEngine {
+impl<D: WorkflowData, S: StateStore<D> + 'static> std::fmt::Debug for WorkflowEngine<D, S> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("WorkflowEngine")
             .field("definitions_count", &self.definitions.read().len())
-            .field("state_count", &self.state_store.count())
             .finish()
     }
 }

sgl-model-gateway/src/workflow/executor.rs

@@ -2,13 +2,13 @@
 
 use async_trait::async_trait;
 
-use super::types::{StepResult, WorkflowContext, WorkflowError, WorkflowResult};
+use super::types::{StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult};
 
 /// Trait for executing individual workflow steps
 #[async_trait]
-pub trait StepExecutor: Send + Sync {
+pub trait StepExecutor<D: WorkflowData>: Send + Sync {
     /// Execute the step with the given context
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult>;
+    async fn execute(&self, context: &mut WorkflowContext<D>) -> WorkflowResult<StepResult>;
 
     /// Check if an error is retry-able
     ///
@@ -22,7 +22,7 @@ pub trait StepExecutor: Send + Sync {
     ///
     /// This hook allows steps to perform cleanup or additional actions
     /// after successful execution.
-    async fn on_success(&self, _context: &WorkflowContext) -> WorkflowResult<()> {
+    async fn on_success(&self, _context: &WorkflowContext<D>) -> WorkflowResult<()> {
         Ok(())
     }
 
@@ -32,7 +32,7 @@ pub trait StepExecutor: Send + Sync {
     /// when the step cannot complete successfully.
     async fn on_failure(
         &self,
-        _context: &WorkflowContext,
+        _context: &WorkflowContext<D>,
         _error: &WorkflowError,
     ) -> WorkflowResult<()> {
         Ok(())
@@ -40,59 +40,82 @@ pub trait StepExecutor: Send + Sync {
 }
 
 /// Simple function-based step executor
-pub struct FunctionStep<F>
+pub struct FunctionStep<D, F>
 where
+    D: WorkflowData,
     F: Fn(
-            &mut WorkflowContext,
+            &mut WorkflowContext<D>,
         ) -> std::pin::Pin<
             Box<dyn std::future::Future<Output = WorkflowResult<StepResult>> + Send + '_>,
         > + Send
         + Sync,
 {
     func: F,
+    _phantom: std::marker::PhantomData<D>,
 }
 
-impl<F> FunctionStep<F>
+impl<D, F> FunctionStep<D, F>
 where
+    D: WorkflowData,
     F: Fn(
-            &mut WorkflowContext,
+            &mut WorkflowContext<D>,
         ) -> std::pin::Pin<
             Box<dyn std::future::Future<Output = WorkflowResult<StepResult>> + Send + '_>,
         > + Send
         + Sync,
 {
     pub fn new(func: F) -> Self {
-        Self { func }
+        Self {
+            func,
+            _phantom: std::marker::PhantomData,
+        }
     }
 }
 
 #[async_trait]
-impl<F> StepExecutor for FunctionStep<F>
+impl<D, F> StepExecutor<D> for FunctionStep<D, F>
 where
+    D: WorkflowData,
     F: Fn(
-            &mut WorkflowContext,
+            &mut WorkflowContext<D>,
         ) -> std::pin::Pin<
             Box<dyn std::future::Future<Output = WorkflowResult<StepResult>> + Send + '_>,
         > + Send
         + Sync,
 {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(&self, context: &mut WorkflowContext<D>) -> WorkflowResult<StepResult> {
         (self.func)(context).await
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use serde::{Deserialize, Serialize};
+
     use super::*;
     use crate::workflow::types::WorkflowInstanceId;
 
+    #[derive(Debug, Clone, Serialize, Deserialize)]
+    struct TestData {
+        value: i32,
+    }
+
+    impl WorkflowData for TestData {
+        fn workflow_type() -> &'static str {
+            "test"
+        }
+    }
+
     struct TestStep {
         should_succeed: bool,
     }
 
     #[async_trait]
-    impl StepExecutor for TestStep {
+    impl StepExecutor<TestData> for TestStep {
-        async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        async fn execute(
+            &self,
+            _context: &mut WorkflowContext<TestData>,
+        ) -> WorkflowResult<StepResult> {
             if self.should_succeed {
                 Ok(StepResult::Success)
             } else {
@@ -109,7 +132,7 @@ mod tests {
         let step = TestStep {
             should_succeed: true,
         };
-        let mut context = WorkflowContext::new(WorkflowInstanceId::new());
+        let mut context = WorkflowContext::new(WorkflowInstanceId::new(), TestData { value: 42 });
 
         let result = step.execute(&mut context).await;
         assert!(result.is_ok());
@@ -121,7 +144,7 @@ mod tests {
         let step = TestStep {
             should_succeed: false,
         };
-        let mut context = WorkflowContext::new(WorkflowInstanceId::new());
+        let mut context = WorkflowContext::new(WorkflowInstanceId::new(), TestData { value: 42 });
 
         let result = step.execute(&mut context).await;
         assert!(result.is_err());

sgl-model-gateway/src/workflow/mod.rs

@@ -11,5 +11,5 @@ pub use definition::{StepDefinition, WorkflowDefinition};
 pub use engine::WorkflowEngine;
 pub use event::{EventBus, EventSubscriber, LoggingSubscriber, WorkflowEvent};
 pub use executor::{FunctionStep, StepExecutor};
-pub use state::WorkflowStateStore;
+pub use state::{InMemoryStore, StateStore};
 pub use types::*;

sgl-model-gateway/src/workflow/state.rs

@@ -1,113 +1,64 @@
 //! Workflow state management
 
-use std::{collections::HashMap, sync::Arc};
+use std::{collections::HashMap, marker::PhantomData, sync::Arc, time::Duration};
 
 use parking_lot::RwLock;
 
 use super::types::{
-    WorkflowError, WorkflowInstanceId, WorkflowResult, WorkflowState, WorkflowStatus,
+    WorkflowContext, WorkflowData, WorkflowError, WorkflowInstanceId, WorkflowResult,
+    WorkflowState, WorkflowStatus,
 };
 
+/// Trait for workflow state persistence.
+///
+/// Implement this trait to provide custom storage backends (e.g., PostgreSQL, Redis).
+/// The default implementation is `InMemoryStore` which keeps state in memory.
+pub trait StateStore<D: WorkflowData>: Send + Sync + Clone {
+    /// Save workflow state
+    fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()>;
+
+    /// Load workflow state by instance ID
+    fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>>;
+
+    /// Update workflow state using a closure
+    fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
+    where
+        F: FnOnce(&mut WorkflowState<D>);
+
+    /// Delete workflow state
+    fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()>;
+
+    /// List all active workflows (Running or Pending)
+    fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>>;
+
+    /// List all workflows
+    fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>>;
+
+    /// Check if workflow is cancelled without loading full state
+    fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool>;
+
+    /// Clean up old completed/failed/cancelled workflows beyond a time threshold
+    fn cleanup_old_workflows(&self, ttl: Duration) -> usize;
+
+    /// Get just the workflow context without cloning the entire state
+    fn get_context(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowContext<D>>;
+}
+
 /// In-memory state storage for workflow instances
 #[derive(Clone)]
-pub struct WorkflowStateStore {
+pub struct InMemoryStore<D: WorkflowData> {
-    states: Arc<RwLock<HashMap<WorkflowInstanceId, WorkflowState>>>,
+    states: Arc<RwLock<HashMap<WorkflowInstanceId, WorkflowState<D>>>>,
+    _phantom: PhantomData<D>,
 }
 
-impl WorkflowStateStore {
+impl<D: WorkflowData> InMemoryStore<D> {
     pub fn new() -> Self {
         Self {
             states: Arc::new(RwLock::new(HashMap::new())),
+            _phantom: PhantomData,
         }
     }
 
-    /// Save workflow state
-    ///
-    /// # Note
-    ///
-    /// This emits a debug log if the workflow context contains unserializable data,
-    /// which would be lost if state persistence is later implemented.
-    pub fn save(&self, state: WorkflowState) -> WorkflowResult<()> {
-        if state.context.has_unserializable_data() {
-            tracing::debug!(
-                instance_id = %state.instance_id,
-                data_count = state.context.data_len(),
-                "Saving workflow state with {} unserializable context entries. \
-                 This data cannot be persisted and will be lost on restart.",
-                state.context.data_len()
-            );
-        }
-        self.states.write().insert(state.instance_id, state);
-        Ok(())
-    }
-
-    /// Load workflow state by instance ID
-    pub fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState> {
-        self.states
-            .read()
-            .get(&instance_id)
-            .cloned()
-            .ok_or(WorkflowError::NotFound(instance_id))
-    }
-
-    /// List all active workflows (Running or Pending)
-    pub fn list_active(&self) -> WorkflowResult<Vec<WorkflowState>> {
-        let states = self.states.read();
-        Ok(states
-            .values()
-            .filter(|s| matches!(s.status, WorkflowStatus::Running | WorkflowStatus::Pending))
-            .cloned()
-            .collect())
-    }
-
-    /// List all workflows
-    pub fn list_all(&self) -> WorkflowResult<Vec<WorkflowState>> {
-        let states = self.states.read();
-        Ok(states.values().cloned().collect())
-    }
-
-    /// Delete workflow state
-    pub fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> {
-        self.states.write().remove(&instance_id);
-        Ok(())
-    }
-
-    /// Update workflow state using a closure
-    pub fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
-    where
-        F: FnOnce(&mut WorkflowState),
-    {
-        let mut states = self.states.write();
-        let state = states
-            .get_mut(&instance_id)
-            .ok_or(WorkflowError::NotFound(instance_id))?;
-        f(state);
-        state.updated_at = chrono::Utc::now();
-        Ok(())
-    }
-
-    /// Get just the workflow context without cloning the entire state.
-    /// More efficient when you only need the context for step execution.
-    pub fn get_context(
-        &self,
-        instance_id: WorkflowInstanceId,
-    ) -> WorkflowResult<super::types::WorkflowContext> {
-        self.states
-            .read()
-            .get(&instance_id)
-            .map(|s| s.context.clone())
-            .ok_or(WorkflowError::NotFound(instance_id))
-    }
-
-    /// Check if workflow is cancelled without loading full state
-    pub fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool> {
-        self.states
-            .read()
-            .get(&instance_id)
-            .map(|s| s.status == WorkflowStatus::Cancelled)
-            .ok_or(WorkflowError::NotFound(instance_id))
-    }
-
     /// Get count of workflows by status
     pub fn count_by_status(&self, status: WorkflowStatus) -> usize {
         self.states
@@ -122,23 +73,91 @@ impl WorkflowStateStore {
         self.states.read().len()
     }
 
-    /// Clean up old completed/failed/cancelled workflows beyond a time threshold
+    /// Clean up a specific completed workflow immediately
-    ///
+    pub fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool {
-    /// This prevents unbounded memory growth by removing workflow states that
+        let mut states = self.states.write();
-    /// have been in a terminal state (Completed, Failed, Cancelled) for longer
+        if let Some(state) = states.get(&instance_id) {
-    /// than the specified TTL (time-to-live).
+            if matches!(
-    ///
+                state.status,
-    /// Active workflows (Running, Pending, Paused) are never cleaned up.
+                WorkflowStatus::Completed | WorkflowStatus::Failed | WorkflowStatus::Cancelled
-    ///
+            ) {
-    /// # Arguments
+                states.remove(&instance_id);
-    ///
+                return true;
-    /// * `ttl` - Time-to-live for terminal workflows. Workflows in terminal states
+            }
-    ///   older than this will be removed.
+        }
-    ///
+        false
-    /// # Returns
+    }
-    ///
+}
-    /// The number of workflow states removed.
+
-    pub fn cleanup_old_workflows(&self, ttl: std::time::Duration) -> usize {
+impl<D: WorkflowData> Default for InMemoryStore<D> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
+    fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()> {
+        self.states.write().insert(state.instance_id, state);
+        Ok(())
+    }
+
+    fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>> {
+        self.states
+            .read()
+            .get(&instance_id)
+            .cloned()
+            .ok_or(WorkflowError::NotFound(instance_id))
+    }
+
+    fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>> {
+        let states = self.states.read();
+        Ok(states
+            .values()
+            .filter(|s| matches!(s.status, WorkflowStatus::Running | WorkflowStatus::Pending))
+            .cloned()
+            .collect())
+    }
+
+    fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>> {
+        let states = self.states.read();
+        Ok(states.values().cloned().collect())
+    }
+
+    fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> {
+        self.states.write().remove(&instance_id);
+        Ok(())
+    }
+
+    fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
+    where
+        F: FnOnce(&mut WorkflowState<D>),
+    {
+        let mut states = self.states.write();
+        let state = states
+            .get_mut(&instance_id)
+            .ok_or(WorkflowError::NotFound(instance_id))?;
+        f(state);
+        state.updated_at = chrono::Utc::now();
+        Ok(())
+    }
+
+    fn get_context(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowContext<D>> {
+        self.states
+            .read()
+            .get(&instance_id)
+            .map(|s| s.context.clone())
+            .ok_or(WorkflowError::NotFound(instance_id))
+    }
+
+    fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool> {
+        self.states
+            .read()
+            .get(&instance_id)
+            .map(|s| s.status == WorkflowStatus::Cancelled)
+            .ok_or(WorkflowError::NotFound(instance_id))
+    }
+
+    fn cleanup_old_workflows(&self, ttl: Duration) -> usize {
         let now = chrono::Utc::now();
         let mut states = self.states.write();
         let initial_count = states.len();
@@ -170,28 +189,4 @@ impl WorkflowStateStore {
         }
         removed_count
     }
-
-    /// Clean up a specific completed workflow immediately
-    ///
-    /// This is useful for cleaning up workflows right after they complete
-    /// when you know they won't be queried again.
-    pub fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool {
-        let mut states = self.states.write();
-        if let Some(state) = states.get(&instance_id) {
-            if matches!(
-                state.status,
-                WorkflowStatus::Completed | WorkflowStatus::Failed | WorkflowStatus::Cancelled
-            ) {
-                states.remove(&instance_id);
-                return true;
-            }
-        }
-        false
-    }
-}
-
-impl Default for WorkflowStateStore {
-    fn default() -> Self {
-        Self::new()
-    }
 }

sgl-model-gateway/src/workflow/types.rs

@@ -1,11 +1,36 @@
 //! Core workflow types and definitions
 
-use std::{collections::HashMap, fmt, sync::Arc, time::Duration};
+use std::{collections::HashMap, fmt, time::Duration};
 
 use chrono::{DateTime, Utc};
-use serde::{Deserialize, Serialize};
+use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use uuid::Uuid;
 
+/// Trait for workflow data that can be passed through workflow steps.
+///
+/// Implementing this trait allows your data type to be used as the typed
+/// context for a workflow. The data must be serializable for state persistence.
+///
+/// # Example
+///
+/// ```ignore
+/// #[derive(Debug, Clone, Serialize, Deserialize)]
+/// pub struct MyWorkflowData {
+///     pub config: MyConfig,
+///     pub result: Option<MyResult>,
+///     #[serde(skip, default)]
+///     pub app_context: Option<Arc<AppContext>>,
+/// }
+///
+/// impl WorkflowData for MyWorkflowData {
+///     fn workflow_type() -> &'static str { "my_workflow" }
+/// }
+/// ```
+pub trait WorkflowData: Serialize + DeserializeOwned + Send + Sync + Clone + 'static {
+    /// Human-readable name for logging and identification
+    fn workflow_type() -> &'static str;
+}
+
 /// Unique identifier for a workflow definition
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub struct WorkflowId(String);
@@ -147,19 +172,23 @@ impl Default for StepState {
 
 /// Workflow instance state
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct WorkflowState {
+#[serde(bound(
+    serialize = "D: Serialize",
+    deserialize = "D: serde::de::DeserializeOwned"
+))]
+pub struct WorkflowState<D: WorkflowData> {
     pub instance_id: WorkflowInstanceId,
     pub definition_id: WorkflowId,
     pub status: WorkflowStatus,
     pub current_step: Option<StepId>,
     pub step_states: HashMap<StepId, StepState>,
-    pub context: WorkflowContext,
+    pub context: WorkflowContext<D>,
     pub created_at: DateTime<Utc>,
     pub updated_at: DateTime<Utc>,
 }
 
-impl WorkflowState {
+impl<D: WorkflowData> WorkflowState<D> {
-    pub fn new(instance_id: WorkflowInstanceId, definition_id: WorkflowId) -> Self {
+    pub fn new(instance_id: WorkflowInstanceId, definition_id: WorkflowId, data: D) -> Self {
         let now = Utc::now();
         Self {
             instance_id,
@@ -167,71 +196,36 @@ impl WorkflowState {
             status: WorkflowStatus::Pending,
             current_step: None,
             step_states: HashMap::new(),
-            context: WorkflowContext::new(instance_id),
+            context: WorkflowContext::new(instance_id, data),
             created_at: now,
             updated_at: now,
         }
     }
 }
 
-/// Shared context passed between workflow steps
+/// Shared context passed between workflow steps.
 ///
-/// # Serialization Warning
+/// The context contains typed workflow data that is fully serializable,
+/// enabling state persistence and workflow recovery.
 ///
-/// The `data` field contains type-erased values that cannot be serialized.
+/// # Type Parameter
-/// This means workflow context is **not preserved** across:
-/// - Process restarts
-/// - State persistence to disk
-/// - Network serialization
 ///
-/// The workflow engine only supports **in-memory execution**. If you need
+/// `D` - The workflow-specific data type implementing `WorkflowData`.
-/// durable workflows, consider implementing a custom serializable context type.
+/// This type holds all the state needed by workflow steps and must be
+/// serializable (except for fields marked with `#[serde(skip)]`).
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct WorkflowContext {
+#[serde(bound(
+    serialize = "D: Serialize",
+    deserialize = "D: serde::de::DeserializeOwned"
+))]
+pub struct WorkflowContext<D: WorkflowData> {
     pub instance_id: WorkflowInstanceId,
-    #[serde(skip)]
+    pub data: D,
-    data: HashMap<String, Arc<dyn std::any::Any + Send + Sync>>,
 }
 
-impl WorkflowContext {
+impl<D: WorkflowData> WorkflowContext<D> {
-    pub fn new(instance_id: WorkflowInstanceId) -> Self {
+    pub fn new(instance_id: WorkflowInstanceId, data: D) -> Self {
-        Self {
+        Self { instance_id, data }
-            instance_id,
-            data: HashMap::new(),
-        }
-    }
-
-    /// Store a value in the context (will be wrapped in Arc)
-    pub fn set<T: Send + Sync + 'static>(&mut self, key: impl Into<String>, value: T) {
-        self.data.insert(key.into(), Arc::new(value));
-    }
-
-    /// Store an Arc directly without double-wrapping
-    pub fn set_arc<T: Send + Sync + 'static>(&mut self, key: impl Into<String>, value: Arc<T>) {
-        self.data.insert(key.into(), value);
-    }
-
-    /// Retrieve a value from the context
-    pub fn get<T: Send + Sync + 'static>(&self, key: &str) -> Option<Arc<T>> {
-        self.data
-            .get(key)
-            .and_then(|v| v.clone().downcast::<T>().ok())
-    }
-
-    /// Retrieve a value from the context, returning an error if not found
-    pub fn get_or_err<T: Send + Sync + 'static>(&self, key: &str) -> Result<Arc<T>, WorkflowError> {
-        self.get::<T>(key)
-            .ok_or_else(|| WorkflowError::ContextValueNotFound(key.to_string()))
-    }
-
-    /// Check if the context has any data that would be lost during serialization
-    pub fn has_unserializable_data(&self) -> bool {
-        !self.data.is_empty()
-    }
-
-    /// Get the number of context entries (useful for debugging)
-    pub fn data_len(&self) -> usize {
-        self.data.len()
     }
 }
 
@@ -270,6 +264,12 @@ pub enum WorkflowError {
     #[error("Context value not found: {0}")]
     ContextValueNotFound(String),
 
+    #[error("Type mismatch: expected {expected}, got {actual}")]
+    TypeMismatch {
+        expected: &'static str,
+        actual: &'static str,
+    },
+
     #[error("Engine is shutting down, not accepting new workflows")]
     ShuttingDown,
 }

sgl-model-gateway/tests/common/mod.rs

@@ -358,12 +358,12 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
 
     // Initialize WorkflowEngine and register workflows
     use smg::{
-        core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
+        core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
         workflow::WorkflowEngine,
     };
     let engine = Arc::new(WorkflowEngine::new());
     engine
-        .register_workflow(create_worker_registration_workflow(&config))
+        .register_workflow(create_local_worker_workflow(&config))
         .expect("worker_registration workflow should be valid");
     engine
         .register_workflow(create_worker_removal_workflow())
@@ -491,12 +491,12 @@ pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc<AppCo
 
     // Initialize WorkflowEngine and register workflows
     use smg::{
-        core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
+        core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
         workflow::WorkflowEngine,
     };
     let engine = Arc::new(WorkflowEngine::new());
     engine
-        .register_workflow(create_worker_registration_workflow(&config))
+        .register_workflow(create_local_worker_workflow(&config))
         .expect("worker_registration workflow should be valid");
     engine
         .register_workflow(create_worker_removal_workflow())
@@ -624,12 +624,12 @@ pub async fn create_test_context_with_mcp_config(
 
     // Initialize WorkflowEngine and register workflows
     use smg::{
-        core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
+        core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
         workflow::WorkflowEngine,
     };
     let engine = Arc::new(WorkflowEngine::new());
     engine
-        .register_workflow(create_worker_registration_workflow(&config))
+        .register_workflow(create_local_worker_workflow(&config))
         .expect("worker_registration workflow should be valid");
     engine
         .register_workflow(create_worker_removal_workflow())

sgl-model-gateway/tests/wasm_test.rs

@@ -109,12 +109,12 @@ async fn create_test_context_with_wasm() -> Arc<AppContext> {
 
     // Initialize WorkflowEngine and register workflows
     use smg::{
-        core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
+        core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
         workflow::WorkflowEngine,
     };
     let engine = Arc::new(WorkflowEngine::new());
     engine
-        .register_workflow(create_worker_registration_workflow(&config))
+        .register_workflow(create_local_worker_workflow(&config))
         .expect("worker_registration workflow should be valid");
     engine
         .register_workflow(create_worker_removal_workflow())
@@ -685,8 +685,11 @@ async fn test_wasm_module_execution() {
 
     // Create workflow context for registration
     use smg::{
-        core::steps::WasmModuleConfigRequest,
+        core::steps::{
-        workflow::{WorkflowContext, WorkflowId, WorkflowInstanceId},
+            workflow_data::{AnyWorkflowData, WasmRegistrationWorkflowData},
+            WasmModuleConfigRequest,
+        },
+        workflow::WorkflowId,
     };
 
     let descriptor = WasmModuleDescriptor {
@@ -700,16 +703,18 @@ async fn test_wasm_module_execution() {
     };
 
     let config_request = WasmModuleConfigRequest { descriptor };
-    let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
+    let workflow_data = AnyWorkflowData::WasmRegistration(WasmRegistrationWorkflowData {
-    workflow_context.set_arc("wasm_module_config", Arc::new(config_request));
+        config: config_request,
-    workflow_context.set_arc("app_context", app_context.clone());
+        wasm_bytes: None,
+        sha256_hash: None,
+        file_size_bytes: None,
+        module_uuid: None,
+        app_context: Some(app_context.clone()),
+    });
 
     // Start workflow
     let instance_id = engine
-        .start_workflow(
+        .start_workflow(WorkflowId::new("wasm_module_registration"), workflow_data)
-            WorkflowId::new("wasm_module_registration"),
-            workflow_context,
-        )
         .await
         .expect("Failed to start workflow");
 
@@ -729,9 +734,9 @@ async fn test_wasm_module_execution() {
 
         match state.status {
             smg::workflow::WorkflowStatus::Completed => {
-                // Extract module UUID from context
+                // Extract module UUID from typed workflow data
-                if let Some(uuid_arc) = state.context.get::<Uuid>("module_uuid") {
+                if let AnyWorkflowData::WasmRegistration(ref data) = state.context.data {
-                    module_uuid = Some(*uuid_arc.as_ref());
+                    module_uuid = data.module_uuid;
                 }
                 break;
             }

sgl-model-gateway/tests/workflow_test.rs

@@ -8,9 +8,25 @@ use std::{
     time::Duration,
 };
 
+use serde::{Deserialize, Serialize};
 use smg::workflow::*;
 use tokio::time::sleep;
 
+/// Test workflow data type for integration tests.
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+struct TestWorkflowData {
+    /// Execution count for tracking step invocations
+    pub execution_count: u32,
+    /// Test key for context sharing tests
+    pub test_key: Option<String>,
+}
+
+impl WorkflowData for TestWorkflowData {
+    fn workflow_type() -> &'static str {
+        "test_workflow"
+    }
+}
+
 // Test step that counts invocations
 struct CountingStep {
     counter: Arc<AtomicU32>,
@@ -18,12 +34,15 @@ struct CountingStep {
 }
 
 #[async_trait::async_trait]
-impl StepExecutor for CountingStep {
+impl StepExecutor<TestWorkflowData> for CountingStep {
-    async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
+        &self,
+        context: &mut WorkflowContext<TestWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
         let count = self.counter.fetch_add(1, Ordering::SeqCst) + 1;
 
         // Store count in context
-        context.set("execution_count", count);
+        context.data.execution_count = count;
 
         if count >= self.should_succeed_after {
             Ok(StepResult::Success)
@@ -40,15 +59,18 @@ impl StepExecutor for CountingStep {
 struct AlwaysSucceedStep;
 
 #[async_trait::async_trait]
-impl StepExecutor for AlwaysSucceedStep {
+impl StepExecutor<TestWorkflowData> for AlwaysSucceedStep {
-    async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
+        &self,
+        _context: &mut WorkflowContext<TestWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
         Ok(StepResult::Success)
     }
 }
 
 #[tokio::test]
 async fn test_simple_workflow_execution() {
-    let engine = WorkflowEngine::new();
+    let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
 
     // Subscribe to events for logging
     engine
@@ -74,7 +96,7 @@ async fn test_simple_workflow_execution() {
 
     // Start workflow
     let instance_id = engine
-        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .start_workflow(workflow_id, TestWorkflowData::default())
         .await
         .unwrap();
 
@@ -89,7 +111,7 @@ async fn test_simple_workflow_execution() {
 
 #[tokio::test]
 async fn test_workflow_with_retry() {
-    let engine = WorkflowEngine::new();
+    let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
     engine
         .event_bus()
         .subscribe(Arc::new(LoggingSubscriber))
@@ -119,7 +141,7 @@ async fn test_workflow_with_retry() {
 
     // Start workflow
     let instance_id = engine
-        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .start_workflow(workflow_id, TestWorkflowData::default())
         .await
         .unwrap();
 
@@ -140,7 +162,7 @@ async fn test_workflow_with_retry() {
 
 #[tokio::test]
 async fn test_workflow_failure_after_max_retries() {
-    let engine = WorkflowEngine::new();
+    let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
     engine
         .event_bus()
         .subscribe(Arc::new(LoggingSubscriber))
@@ -170,7 +192,7 @@ async fn test_workflow_failure_after_max_retries() {
 
     // Start workflow
     let instance_id = engine
-        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .start_workflow(workflow_id, TestWorkflowData::default())
         .await
         .unwrap();
 
@@ -191,7 +213,7 @@ async fn test_workflow_failure_after_max_retries() {
 
 #[tokio::test]
 async fn test_workflow_continue_on_failure() {
-    let engine = WorkflowEngine::new();
+    let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
     engine
         .event_bus()
         .subscribe(Arc::new(LoggingSubscriber))
@@ -227,7 +249,7 @@ async fn test_workflow_continue_on_failure() {
 
     // Start workflow
     let instance_id = engine
-        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .start_workflow(workflow_id, TestWorkflowData::default())
         .await
         .unwrap();
 
@@ -249,34 +271,40 @@ async fn test_workflow_continue_on_failure() {
 
 #[tokio::test]
 async fn test_workflow_context_sharing() {
-    let engine = WorkflowEngine::new();
+    let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
 
     struct ContextWriterStep {
-        key: String,
         value: String,
     }
 
     #[async_trait::async_trait]
-    impl StepExecutor for ContextWriterStep {
+    impl StepExecutor<TestWorkflowData> for ContextWriterStep {
-        async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        async fn execute(
-            context.set(self.key.clone(), self.value.clone());
+            &self,
+            context: &mut WorkflowContext<TestWorkflowData>,
+        ) -> WorkflowResult<StepResult> {
+            context.data.test_key = Some(self.value.clone());
             Ok(StepResult::Success)
         }
     }
 
     struct ContextReaderStep {
-        key: String,
         expected_value: String,
     }
 
     #[async_trait::async_trait]
-    impl StepExecutor for ContextReaderStep {
+    impl StepExecutor<TestWorkflowData> for ContextReaderStep {
-        async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        async fn execute(
-            let value: Arc<String> = context
+            &self,
-                .get(&self.key)
+            context: &mut WorkflowContext<TestWorkflowData>,
-                .ok_or_else(|| WorkflowError::ContextValueNotFound(self.key.clone()))?;
+        ) -> WorkflowResult<StepResult> {
+            let value = context
+                .data
+                .test_key
+                .as_ref()
+                .ok_or_else(|| WorkflowError::ContextValueNotFound("test_key".to_string()))?;
 
-            if *value == self.expected_value {
+            if value == &self.expected_value {
                 Ok(StepResult::Success)
             } else {
                 Err(WorkflowError::StepFailed {
@@ -292,7 +320,6 @@ async fn test_workflow_context_sharing() {
             "writer",
             "Write to context",
             Arc::new(ContextWriterStep {
-                key: "test_key".to_string(),
                 value: "test_value".to_string(),
             }),
         ))
@@ -300,7 +327,6 @@ async fn test_workflow_context_sharing() {
             "reader",
             "Read from context",
             Arc::new(ContextReaderStep {
-                key: "test_key".to_string(),
                 expected_value: "test_value".to_string(),
             }),
         ));
@@ -309,7 +335,7 @@ async fn test_workflow_context_sharing() {
     engine.register_workflow(workflow).unwrap();
 
     let instance_id = engine
-        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .start_workflow(workflow_id, TestWorkflowData::default())
         .await
         .unwrap();
 
@@ -332,8 +358,11 @@ struct TimingStep {
 }
 
 #[async_trait::async_trait]
-impl StepExecutor for TimingStep {
+impl StepExecutor<TestWorkflowData> for TimingStep {
-    async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+    async fn execute(
+        &self,
+        _context: &mut WorkflowContext<TestWorkflowData>,
+    ) -> WorkflowResult<StepResult> {
         let start = std::time::Instant::now();
         self.start_times
             .write()
@@ -351,7 +380,7 @@ impl StepExecutor for TimingStep {
 #[tokio::test]
 async fn test_parallel_execution_no_dependencies() {
     // Steps without dependencies should run in parallel
-    let engine = WorkflowEngine::new();
+    let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
 
     let start_times: Arc<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
         Arc::new(parking_lot::RwLock::new(Vec::new()));
@@ -398,7 +427,7 @@ async fn test_parallel_execution_no_dependencies() {
 
     let overall_start = std::time::Instant::now();
     let instance_id = engine
-        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .start_workflow(workflow_id, TestWorkflowData::default())
         .await
         .unwrap();
 
@@ -446,7 +475,7 @@ async fn test_dag_with_dependencies() {
     //   A ──┐
     //       ├──> C
     //   B ──┘
-    let engine = WorkflowEngine::new();
+    let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
 
     let start_times: Arc<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
         Arc::new(parking_lot::RwLock::new(Vec::new()));
@@ -492,7 +521,7 @@ async fn test_dag_with_dependencies() {
     engine.register_workflow(workflow).unwrap();
 
     let instance_id = engine
-        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .start_workflow(workflow_id, TestWorkflowData::default())
         .await
         .unwrap();
 
@@ -523,15 +552,18 @@ async fn test_dag_with_dependencies() {
 #[tokio::test]
 async fn test_dag_dependency_failure_blocks_dependents() {
     // If step A fails with FailWorkflow, step B (depends on A) should not run
-    let engine = WorkflowEngine::new();
+    let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
 
     let b_executed = Arc::new(AtomicU32::new(0));
 
     struct FailingStep;
 
     #[async_trait::async_trait]
-    impl StepExecutor for FailingStep {
+    impl StepExecutor<TestWorkflowData> for FailingStep {
-        async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        async fn execute(
+            &self,
+            _context: &mut WorkflowContext<TestWorkflowData>,
+        ) -> WorkflowResult<StepResult> {
             Err(WorkflowError::StepFailed {
                 step_id: StepId::new("failing"),
                 message: "Intentional failure".to_string(),
@@ -548,8 +580,11 @@ async fn test_dag_dependency_failure_blocks_dependents() {
     }
 
     #[async_trait::async_trait]
-    impl StepExecutor for TrackingStep {
+    impl StepExecutor<TestWorkflowData> for TrackingStep {
-        async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        async fn execute(
+            &self,
+            _context: &mut WorkflowContext<TestWorkflowData>,
+        ) -> WorkflowResult<StepResult> {
             self.counter.fetch_add(1, Ordering::SeqCst);
             Ok(StepResult::Success)
         }
@@ -575,7 +610,7 @@ async fn test_dag_dependency_failure_blocks_dependents() {
     engine.register_workflow(workflow).unwrap();
 
     let instance_id = engine
-        .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
+        .start_workflow(workflow_id, TestWorkflowData::default())
         .await
         .unwrap();