[model-gateway] refactor workflow engine from type erasure to typed engines (#16973)

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

@@ -781,6 +781,65 @@ impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
         self.state_store.load(instance_id)
     }
 
+    /// Wait for a workflow to complete with adaptive polling
+    ///
+    /// Returns Ok with success message on completion, Err on failure/timeout/cancellation.
+    /// Automatically cleans up terminal workflow states.
+    pub async fn wait_for_completion(
+        &self,
+        instance_id: WorkflowInstanceId,
+        label: &str,
+        timeout_duration: Duration,
+    ) -> Result<String, String> {
+        let start = std::time::Instant::now();
+        let mut poll_interval = Duration::from_millis(100);
+        let max_poll_interval = Duration::from_millis(2000);
+        let poll_backoff = Duration::from_millis(200);
+
+        loop {
+            if start.elapsed() > timeout_duration {
+                return Err(format!(
+                    "Workflow timeout after {}s for {}",
+                    timeout_duration.as_secs(),
+                    label
+                ));
+            }
+
+            let state = self
+                .get_status(instance_id)
+                .map_err(|e| format!("Failed to get workflow status: {:?}", e))?;
+
+            let result = match state.status {
+                WorkflowStatus::Completed => {
+                    Ok(format!("{} completed successfully via workflow", label))
+                }
+                WorkflowStatus::Failed => {
+                    let current_step = state.current_step.as_ref();
+                    let step_name = current_step
+                        .map(|s| s.to_string())
+                        .unwrap_or_else(|| "unknown".to_string());
+                    let error_msg = current_step
+                        .and_then(|step_id| state.step_states.get(step_id))
+                        .and_then(|s| s.last_error.as_deref())
+                        .unwrap_or("Unknown error");
+                    Err(format!(
+                        "Workflow failed at step {}: {}",
+                        step_name, error_msg
+                    ))
+                }
+                WorkflowStatus::Cancelled => Err(format!("Workflow cancelled for {}", label)),
+                WorkflowStatus::Pending | WorkflowStatus::Paused | WorkflowStatus::Running => {
+                    tokio::time::sleep(poll_interval).await;
+                    poll_interval = (poll_interval + poll_backoff).min(max_poll_interval);
+                    continue;
+                }
+            };
+
+            self.state_store.cleanup_if_terminal(instance_id);
+            return result;
+        }
+    }
+
     /// Clone engine for async execution
     fn clone_for_execution(&self) -> Self {
         Self {

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

@@ -42,6 +42,10 @@ pub trait StateStore<D: WorkflowData>: Send + Sync + Clone {
 
     /// Get just the workflow context without cloning the entire state
     fn get_context(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowContext<D>>;
+
+    /// Clean up a specific workflow immediately if it's in a terminal state
+    /// Returns true if the workflow was removed, false otherwise
+    fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool;
 }
 
 /// In-memory state storage for workflow instances
@@ -72,21 +76,6 @@ impl<D: WorkflowData> InMemoryStore<D> {
     pub fn count(&self) -> usize {
         self.states.read().len()
     }
-
-    /// Clean up a specific completed workflow immediately
-    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<D: WorkflowData> Default for InMemoryStore<D> {
@@ -189,4 +178,18 @@ impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
         }
         removed_count
     }
+
+    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
+    }
 }