From 10c68f6236075699707527c047340d81ce764639 Mon Sep 17 00:00:00 2001
From: Simo Lin <linsimo.mark@gmail.com>
Date: Fri, 12 Dec 2025 06:29:12 -0800
Subject: [PATCH] [model-gateway] refactor: extract workflow engine to
 src/workflow module (#14996)

---
 sgl-model-gateway/src/app_context.rs           |  3 ++-
 sgl-model-gateway/src/core/job_queue.rs        | 10 ++++------
 sgl-model-gateway/src/core/mod.rs              |  4 ++--
 .../steps/external_worker_registration.rs      |  2 +-
 .../{workflow => }/steps/mcp_registration.rs   |  2 +-
 .../src/core/{workflow => }/steps/mod.rs       |  0
 .../steps/wasm_module_registration.rs          |  2 +-
 .../steps/wasm_module_removal.rs               |  2 +-
 .../steps/worker_registration.rs               |  5 +++--
 .../{workflow => }/steps/worker_removal.rs     |  5 +----
 sgl-model-gateway/src/lib.rs                   |  1 +
 sgl-model-gateway/src/server.rs                |  9 ++++-----
 sgl-model-gateway/src/wasm/route.rs            |  4 ++--
 .../src/{core => }/workflow/definition.rs      |  0
 .../src/{core => }/workflow/engine.rs          |  0
 .../src/{core => }/workflow/event.rs           |  0
 .../src/{core => }/workflow/executor.rs        |  4 ++--
 .../src/{core => }/workflow/mod.rs             |  6 ------
 .../src/{core => }/workflow/state.rs           |  0
 .../src/{core => }/workflow/types.rs           |  0
 sgl-model-gateway/tests/common/mod.rs          | 10 ++++++----
 sgl-model-gateway/tests/wasm_test.rs           | 18 +++++++++---------
 sgl-model-gateway/tests/workflow_test.rs       |  2 +-
 23 files changed, 41 insertions(+), 48 deletions(-)
 rename sgl-model-gateway/src/core/{workflow => }/steps/external_worker_registration.rs (99%)
 rename sgl-model-gateway/src/core/{workflow => }/steps/mcp_registration.rs (99%)
 rename sgl-model-gateway/src/core/{workflow => }/steps/mod.rs (100%)
 rename sgl-model-gateway/src/core/{workflow => }/steps/wasm_module_registration.rs (99%)
 rename sgl-model-gateway/src/core/{workflow => }/steps/wasm_module_removal.rs (98%)
 rename sgl-model-gateway/src/core/{workflow => }/steps/worker_registration.rs (99%)
 rename sgl-model-gateway/src/core/{workflow => }/steps/worker_removal.rs (99%)
 rename sgl-model-gateway/src/{core => }/workflow/definition.rs (100%)
 rename sgl-model-gateway/src/{core => }/workflow/engine.rs (100%)
 rename sgl-model-gateway/src/{core => }/workflow/event.rs (100%)
 rename sgl-model-gateway/src/{core => }/workflow/executor.rs (96%)
 rename sgl-model-gateway/src/{core => }/workflow/mod.rs (60%)
 rename sgl-model-gateway/src/{core => }/workflow/state.rs (100%)
 rename sgl-model-gateway/src/{core => }/workflow/types.rs (100%)

diff --git a/sgl-model-gateway/src/app_context.rs b/sgl-model-gateway/src/app_context.rs
index 2eb82ce6d..d14567c47 100644
--- a/sgl-model-gateway/src/app_context.rs
+++ b/sgl-model-gateway/src/app_context.rs
@@ -8,7 +8,7 @@ use tracing::info;
 
 use crate::{
     config::RouterConfig,
-    core::{workflow::WorkflowEngine, ConnectionMode, JobQueue, LoadMonitor, WorkerRegistry},
+    core::{ConnectionMode, JobQueue, LoadMonitor, WorkerRegistry},
     data_connector::{
         create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage,
     },
@@ -24,6 +24,7 @@ use crate::{
     },
     tool_parser::ParserFactory as ToolParserFactory,
     wasm::{config::WasmRuntimeConfig, module_manager::WasmModuleManager},
+    workflow::WorkflowEngine,
 };
 
 /// Error type for AppContext builder
diff --git a/sgl-model-gateway/src/core/job_queue.rs b/sgl-model-gateway/src/core/job_queue.rs
index 68c633fe7..fa2385661 100644
--- a/sgl-model-gateway/src/core/job_queue.rs
+++ b/sgl-model-gateway/src/core/job_queue.rs
@@ -16,16 +16,14 @@ use tracing::{debug, error, info, warn};
 use crate::{
     app_context::AppContext,
     config::{RouterConfig, RoutingMode},
-    core::workflow::{
-        steps::{
-            McpServerConfigRequest, WasmModuleConfigRequest, WasmModuleRemovalRequest,
-            WorkerRemovalRequest,
-        },
-        WorkflowContext, WorkflowEngine, WorkflowId, WorkflowInstanceId, WorkflowStatus,
+    core::steps::{
+        McpServerConfigRequest, WasmModuleConfigRequest, WasmModuleRemovalRequest,
+        WorkerRemovalRequest,
     },
     mcp::McpConfig,
     observability::metrics::RouterMetrics,
     protocols::worker_spec::{JobStatus, WorkerConfigRequest},
+    workflow::{WorkflowContext, WorkflowEngine, WorkflowId, WorkflowInstanceId, WorkflowStatus},
 };
 
 /// Job types for control plane operations
diff --git a/sgl-model-gateway/src/core/mod.rs b/sgl-model-gateway/src/core/mod.rs
index ab7191d8e..8e0f7db28 100644
--- a/sgl-model-gateway/src/core/mod.rs
+++ b/sgl-model-gateway/src/core/mod.rs
@@ -6,7 +6,7 @@
 //! - Error types
 //! - Circuit breaker for reliability
 //! - Token buckets for rate limiting
-//! - Workflow engine for multi-step operations
+//! - Workflow steps for multi-step operations
 //! - Common utilities
 
 pub mod circuit_breaker;
@@ -16,12 +16,12 @@ pub mod metrics_aggregator;
 pub mod model_card;
 pub mod model_type;
 pub mod retry;
+pub mod steps;
 pub mod token_bucket;
 pub mod worker;
 pub mod worker_builder;
 pub mod worker_manager;
 pub mod worker_registry;
-pub mod workflow;
 
 pub use circuit_breaker::{
     CircuitBreaker, CircuitBreakerConfig, CircuitBreakerStats, CircuitState,
diff --git a/sgl-model-gateway/src/core/workflow/steps/external_worker_registration.rs b/sgl-model-gateway/src/core/steps/external_worker_registration.rs
similarity index 99%
rename from sgl-model-gateway/src/core/workflow/steps/external_worker_registration.rs
rename to sgl-model-gateway/src/core/steps/external_worker_registration.rs
index cd54067dd..101b1abef 100644
--- a/sgl-model-gateway/src/core/workflow/steps/external_worker_registration.rs
+++ b/sgl-model-gateway/src/core/steps/external_worker_registration.rs
@@ -28,11 +28,11 @@ use crate::{
     core::{
         model_card::{ModelCard, ProviderType},
         model_type::ModelType,
-        workflow::*,
         BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, HealthConfig, RuntimeType,
         Worker, WorkerType,
     },
     protocols::worker_spec::WorkerConfigRequest,
+    workflow::*,
 };
 
 // HTTP client for API calls
diff --git a/sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs b/sgl-model-gateway/src/core/steps/mcp_registration.rs
similarity index 99%
rename from sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs
rename to sgl-model-gateway/src/core/steps/mcp_registration.rs
index f7024bead..f236a80a2 100644
--- a/sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs
+++ b/sgl-model-gateway/src/core/steps/mcp_registration.rs
@@ -16,8 +16,8 @@ use tracing::{debug, error, info, warn};
 
 use crate::{
     app_context::AppContext,
-    core::workflow::*,
     mcp::{config::McpServerConfig, manager::McpManager},
+    workflow::*,
 };
 
 /// MCP server connection configuration
diff --git a/sgl-model-gateway/src/core/workflow/steps/mod.rs b/sgl-model-gateway/src/core/steps/mod.rs
similarity index 100%
rename from sgl-model-gateway/src/core/workflow/steps/mod.rs
rename to sgl-model-gateway/src/core/steps/mod.rs
diff --git a/sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs b/sgl-model-gateway/src/core/steps/wasm_module_registration.rs
similarity index 99%
rename from sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs
rename to sgl-model-gateway/src/core/steps/wasm_module_registration.rs
index 8345c0d12..417d5bad3 100644
--- a/sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs
+++ b/sgl-model-gateway/src/core/steps/wasm_module_registration.rs
@@ -24,8 +24,8 @@ use wasmtime::{component::Component, Config, Engine};
 
 use crate::{
     app_context::AppContext,
-    core::workflow::*,
     wasm::module::{WasmModule, WasmModuleDescriptor, WasmModuleMeta},
+    workflow::*,
 };
 
 /// WASM module registration request
diff --git a/sgl-model-gateway/src/core/workflow/steps/wasm_module_removal.rs b/sgl-model-gateway/src/core/steps/wasm_module_removal.rs
similarity index 98%
rename from sgl-model-gateway/src/core/workflow/steps/wasm_module_removal.rs
rename to sgl-model-gateway/src/core/steps/wasm_module_removal.rs
index 9cd20c9d2..3a0efc023 100644
--- a/sgl-model-gateway/src/core/workflow/steps/wasm_module_removal.rs
+++ b/sgl-model-gateway/src/core/steps/wasm_module_removal.rs
@@ -12,7 +12,7 @@ use async_trait::async_trait;
 use tracing::{debug, info};
 use uuid::Uuid;
 
-use crate::{app_context::AppContext, core::workflow::*};
+use crate::{app_context::AppContext, workflow::*};
 
 /// WASM module removal request
 #[derive(Debug, Clone)]
diff --git a/sgl-model-gateway/src/core/workflow/steps/worker_registration.rs b/sgl-model-gateway/src/core/steps/worker_registration.rs
similarity index 99%
rename from sgl-model-gateway/src/core/workflow/steps/worker_registration.rs
rename to sgl-model-gateway/src/core/steps/worker_registration.rs
index b7bacbb20..d4d4dee8d 100644
--- a/sgl-model-gateway/src/core/workflow/steps/worker_registration.rs
+++ b/sgl-model-gateway/src/core/steps/worker_registration.rs
@@ -24,11 +24,12 @@ use tracing::{debug, info, warn};
 use crate::{
     app_context::AppContext,
     core::{
-        workflow::*, BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode,
-        DPAwareWorkerBuilder, HealthConfig, ModelCard, RuntimeType, Worker, WorkerType,
+        BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, DPAwareWorkerBuilder,
+        HealthConfig, ModelCard, RuntimeType, Worker, WorkerType,
     },
     protocols::worker_spec::WorkerConfigRequest,
     routers::grpc::client::GrpcClient,
+    workflow::*,
 };
 
 // HTTP client for metadata fetching
diff --git a/sgl-model-gateway/src/core/workflow/steps/worker_removal.rs b/sgl-model-gateway/src/core/steps/worker_removal.rs
similarity index 99%
rename from sgl-model-gateway/src/core/workflow/steps/worker_removal.rs
rename to sgl-model-gateway/src/core/steps/worker_removal.rs
index 325029614..f8770ab7e 100644
--- a/sgl-model-gateway/src/core/workflow/steps/worker_removal.rs
+++ b/sgl-model-gateway/src/core/steps/worker_removal.rs
@@ -14,10 +14,7 @@ use std::{collections::HashSet, sync::Arc, time::Duration};
 use async_trait::async_trait;
 use tracing::{debug, info};
 
-use crate::{
-    app_context::AppContext,
-    core::{workflow::*, Worker},
-};
+use crate::{app_context::AppContext, core::Worker, workflow::*};
 
 /// Request structure for worker removal
 #[derive(Debug, Clone)]
diff --git a/sgl-model-gateway/src/lib.rs b/sgl-model-gateway/src/lib.rs
index 91b02976c..dac820881 100644
--- a/sgl-model-gateway/src/lib.rs
+++ b/sgl-model-gateway/src/lib.rs
@@ -17,3 +17,4 @@ pub mod tokenizer;
 pub mod tool_parser;
 pub mod version;
 pub mod wasm;
+pub mod workflow;
diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs
index f3dd82b32..8cc03b85a 100644
--- a/sgl-model-gateway/src/server.rs
+++ b/sgl-model-gateway/src/server.rs
@@ -23,14 +23,12 @@ use crate::{
     app_context::AppContext,
     config::{RouterConfig, RoutingMode},
     core::{
-        worker_to_info,
-        workflow::{
+        steps::{
             create_external_worker_registration_workflow, create_mcp_registration_workflow,
             create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
-            create_worker_registration_workflow, create_worker_removal_workflow, LoggingSubscriber,
-            WorkflowEngine,
+            create_worker_registration_workflow, create_worker_removal_workflow,
         },
-        Job, JobQueue, JobQueueConfig, WorkerManager, WorkerType,
+        worker_to_info, Job, JobQueue, JobQueueConfig, WorkerManager, WorkerType,
     },
     middleware::{self, AuthConfig, QueuedRequest},
     observability::{
@@ -52,6 +50,7 @@ use crate::{
     routers::{conversations, router_manager::RouterManager, RouterTrait},
     service_discovery::{start_service_discovery, ServiceDiscoveryConfig},
     wasm::route::{add_wasm_module, list_wasm_modules, remove_wasm_module},
+    workflow::{LoggingSubscriber, WorkflowEngine},
 };
 
 #[derive(Clone)]
diff --git a/sgl-model-gateway/src/wasm/route.rs b/sgl-model-gateway/src/wasm/route.rs
index 6cbaf09d7..6d9b18d68 100644
--- a/sgl-model-gateway/src/wasm/route.rs
+++ b/sgl-model-gateway/src/wasm/route.rs
@@ -15,7 +15,7 @@ use axum::{
 use uuid::Uuid;
 
 use crate::{
-    core::{job_queue::Job, workflow::steps::WasmModuleConfigRequest},
+    core::{job_queue::Job, steps::WasmModuleConfigRequest},
     server::AppState,
     wasm::module::{
         WasmMetrics, WasmModuleAddRequest, WasmModuleAddResponse, WasmModuleAddResult,
@@ -172,7 +172,7 @@ pub async fn remove_wasm_module(
         return StatusCode::INTERNAL_SERVER_ERROR.into_response();
     };
 
-    use crate::core::workflow::steps::WasmModuleRemovalRequest;
+    use crate::core::steps::WasmModuleRemovalRequest;
 
     let removal_request = WasmModuleRemovalRequest::new(module_uuid);
 
diff --git a/sgl-model-gateway/src/core/workflow/definition.rs b/sgl-model-gateway/src/workflow/definition.rs
similarity index 100%
rename from sgl-model-gateway/src/core/workflow/definition.rs
rename to sgl-model-gateway/src/workflow/definition.rs
diff --git a/sgl-model-gateway/src/core/workflow/engine.rs b/sgl-model-gateway/src/workflow/engine.rs
similarity index 100%
rename from sgl-model-gateway/src/core/workflow/engine.rs
rename to sgl-model-gateway/src/workflow/engine.rs
diff --git a/sgl-model-gateway/src/core/workflow/event.rs b/sgl-model-gateway/src/workflow/event.rs
similarity index 100%
rename from sgl-model-gateway/src/core/workflow/event.rs
rename to sgl-model-gateway/src/workflow/event.rs
diff --git a/sgl-model-gateway/src/core/workflow/executor.rs b/sgl-model-gateway/src/workflow/executor.rs
similarity index 96%
rename from sgl-model-gateway/src/core/workflow/executor.rs
rename to sgl-model-gateway/src/workflow/executor.rs
index 1a27c1922..0c9aa9a6a 100644
--- a/sgl-model-gateway/src/core/workflow/executor.rs
+++ b/sgl-model-gateway/src/workflow/executor.rs
@@ -84,7 +84,7 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::core::workflow::types::WorkflowInstanceId;
+    use crate::workflow::types::WorkflowInstanceId;
 
     struct TestStep {
         should_succeed: bool,
@@ -97,7 +97,7 @@ mod tests {
                 Ok(StepResult::Success)
             } else {
                 Err(WorkflowError::StepFailed {
-                    step_id: crate::core::workflow::types::StepId::new("test"),
+                    step_id: crate::workflow::types::StepId::new("test"),
                     message: "test error".to_string(),
                 })
             }
diff --git a/sgl-model-gateway/src/core/workflow/mod.rs b/sgl-model-gateway/src/workflow/mod.rs
similarity index 60%
rename from sgl-model-gateway/src/core/workflow/mod.rs
rename to sgl-model-gateway/src/workflow/mod.rs
index 9a35ce76e..133bdde60 100644
--- a/sgl-model-gateway/src/core/workflow/mod.rs
+++ b/sgl-model-gateway/src/workflow/mod.rs
@@ -5,7 +5,6 @@ mod engine;
 mod event;
 mod executor;
 mod state;
-pub mod steps;
 pub mod types;
 
 // Re-export main types
@@ -14,9 +13,4 @@ pub use engine::WorkflowEngine;
 pub use event::{EventBus, EventSubscriber, LoggingSubscriber, WorkflowEvent};
 pub use executor::{FunctionStep, StepExecutor};
 pub use state::WorkflowStateStore;
-pub use steps::{
-    create_external_worker_registration_workflow, create_mcp_registration_workflow,
-    create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
-    create_worker_registration_workflow, create_worker_removal_workflow,
-};
 pub use types::*;
diff --git a/sgl-model-gateway/src/core/workflow/state.rs b/sgl-model-gateway/src/workflow/state.rs
similarity index 100%
rename from sgl-model-gateway/src/core/workflow/state.rs
rename to sgl-model-gateway/src/workflow/state.rs
diff --git a/sgl-model-gateway/src/core/workflow/types.rs b/sgl-model-gateway/src/workflow/types.rs
similarity index 100%
rename from sgl-model-gateway/src/core/workflow/types.rs
rename to sgl-model-gateway/src/workflow/types.rs
diff --git a/sgl-model-gateway/tests/common/mod.rs b/sgl-model-gateway/tests/common/mod.rs
index 1a8b01a9d..9f8a8d9c5 100644
--- a/sgl-model-gateway/tests/common/mod.rs
+++ b/sgl-model-gateway/tests/common/mod.rs
@@ -102,8 +102,9 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
         .expect("JobQueue should only be initialized once");
 
     // Initialize WorkflowEngine and register workflows
-    use sgl_model_gateway::core::workflow::{
-        create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine,
+    use sgl_model_gateway::{
+        core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
+        workflow::WorkflowEngine,
     };
     let engine = Arc::new(WorkflowEngine::new());
     engine.register_workflow(create_worker_registration_workflow(&config));
@@ -233,8 +234,9 @@ pub async fn create_test_context_with_mcp_config(
         .expect("JobQueue should only be initialized once");
 
     // Initialize WorkflowEngine and register workflows
-    use sgl_model_gateway::core::workflow::{
-        create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine,
+    use sgl_model_gateway::{
+        core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
+        workflow::WorkflowEngine,
     };
     let engine = Arc::new(WorkflowEngine::new());
     engine.register_workflow(create_worker_registration_workflow(&config));
diff --git a/sgl-model-gateway/tests/wasm_test.rs b/sgl-model-gateway/tests/wasm_test.rs
index bdd4e0060..caccd812b 100644
--- a/sgl-model-gateway/tests/wasm_test.rs
+++ b/sgl-model-gateway/tests/wasm_test.rs
@@ -18,9 +18,7 @@ use axum::{
 use sgl_model_gateway::{
     app_context::AppContext,
     config::RouterConfig,
-    core::workflow::{
-        create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
-    },
+    core::steps::{create_wasm_module_registration_workflow, create_wasm_module_removal_workflow},
     routers::RouterFactory,
     server::{build_app, AppState},
     wasm::{
@@ -113,8 +111,9 @@ async fn create_test_context_with_wasm() -> Arc<AppContext> {
         .expect("JobQueue should only be initialized once");
 
     // Initialize WorkflowEngine and register workflows
-    use sgl_model_gateway::core::workflow::{
-        create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine,
+    use sgl_model_gateway::{
+        core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
+        workflow::WorkflowEngine,
     };
     let engine = Arc::new(WorkflowEngine::new());
     engine.register_workflow(create_worker_registration_workflow(&config));
@@ -684,8 +683,9 @@ async fn test_wasm_module_execution() {
         .expect("Workflow engine should be initialized");
 
     // Create workflow context for registration
-    use sgl_model_gateway::core::workflow::{
-        steps::WasmModuleConfigRequest, WorkflowContext, WorkflowId, WorkflowInstanceId,
+    use sgl_model_gateway::{
+        core::steps::WasmModuleConfigRequest,
+        workflow::{WorkflowContext, WorkflowId, WorkflowInstanceId},
     };
 
     let descriptor = WasmModuleDescriptor {
@@ -727,14 +727,14 @@ async fn test_wasm_module_execution() {
             .expect("Failed to get workflow status");
 
         match state.status {
-            sgl_model_gateway::core::workflow::WorkflowStatus::Completed => {
+            sgl_model_gateway::workflow::WorkflowStatus::Completed => {
                 // Extract module UUID from context
                 if let Some(uuid_arc) = state.context.get::<Uuid>("module_uuid") {
                     module_uuid = Some(*uuid_arc.as_ref());
                 }
                 break;
             }
-            sgl_model_gateway::core::workflow::WorkflowStatus::Failed => {
+            sgl_model_gateway::workflow::WorkflowStatus::Failed => {
                 panic!("Workflow failed: {:?}", state);
             }
             _ => {
diff --git a/sgl-model-gateway/tests/workflow_test.rs b/sgl-model-gateway/tests/workflow_test.rs
index be36030c3..de5c6ca97 100644
--- a/sgl-model-gateway/tests/workflow_test.rs
+++ b/sgl-model-gateway/tests/workflow_test.rs
@@ -8,7 +8,7 @@ use std::{
     time::Duration,
 };
 
-use sgl_model_gateway::core::workflow::*;
+use sgl_model_gateway::workflow::*;
 use tokio::time::sleep;
 
 // Test step that counts invocations