diff --git a/sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs b/sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs index 7952e37c2..fd386f43d 100644 --- a/sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs +++ b/sgl-model-gateway/src/core/workflow/steps/mcp_registration.rs @@ -46,7 +46,8 @@ pub struct ConnectMcpServerStep; #[async_trait] impl StepExecutor for ConnectMcpServerStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context.get_or_err("mcp_server_config")?; + let config_request: Arc = + context.get_or_err("mcp_server_config")?; let app_context: Arc = context.get_or_err("app_context")?; debug!("Connecting to MCP server: {}", config_request.name); @@ -98,7 +99,8 @@ impl StepExecutor for DiscoverMcpInventoryStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { use rmcp::{service::RunningService, RoleClient}; - let config_request: Arc = context.get_or_err("mcp_server_config")?; + let config_request: Arc = + context.get_or_err("mcp_server_config")?; let app_context: Arc = context.get_or_err("app_context")?; let mcp_client: Arc> = context.get_or_err("mcp_client")?; @@ -143,7 +145,8 @@ impl StepExecutor for RegisterMcpServerStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { use rmcp::{service::RunningService, RoleClient}; - let config_request: Arc = context.get_or_err("mcp_server_config")?; + let config_request: Arc = + context.get_or_err("mcp_server_config")?; let app_context: Arc = context.get_or_err("app_context")?; let mcp_client: Arc> = context.get_or_err("mcp_client")?; @@ -183,7 +186,8 @@ pub struct ValidateRegistrationStep; #[async_trait] impl StepExecutor for ValidateRegistrationStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context.get_or_err("mcp_server_config")?; + let config_request: Arc = + context.get_or_err("mcp_server_config")?; let client_registered = context .get::>("mcp_client") diff --git a/sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs b/sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs index 2ee3082c2..8345c0d12 100644 --- a/sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs +++ b/sgl-model-gateway/src/core/workflow/steps/wasm_module_registration.rs @@ -75,7 +75,8 @@ pub struct ValidateDescriptorStep; #[async_trait] impl StepExecutor for ValidateDescriptorStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context.get_or_err("wasm_module_config")?; + let config_request: Arc = + context.get_or_err("wasm_module_config")?; let descriptor = &config_request.descriptor; @@ -233,7 +234,8 @@ pub struct CalculateHashStep; #[async_trait] impl StepExecutor for CalculateHashStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context.get_or_err("wasm_module_config")?; + let config_request: Arc = + context.get_or_err("wasm_module_config")?; let file_path = &config_request.descriptor.file_path; @@ -291,7 +293,8 @@ pub struct CheckDuplicateStep; #[async_trait] impl StepExecutor for CheckDuplicateStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context.get_or_err("wasm_module_config")?; + let config_request: Arc = + context.get_or_err("wasm_module_config")?; let app_context: Arc = context.get_or_err("app_context")?; let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?; @@ -339,7 +342,8 @@ pub struct LoadWasmBytesStep; #[async_trait] impl StepExecutor for LoadWasmBytesStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context.get_or_err("wasm_module_config")?; + let config_request: Arc = + context.get_or_err("wasm_module_config")?; let file_path = &config_request.descriptor.file_path; @@ -374,7 +378,8 @@ pub struct ValidateWasmComponentStep; #[async_trait] impl StepExecutor for ValidateWasmComponentStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context.get_or_err("wasm_module_config")?; + let config_request: Arc = + context.get_or_err("wasm_module_config")?; let wasm_bytes: Arc> = context.get_or_err("wasm_bytes")?; debug!( @@ -425,7 +430,8 @@ pub struct RegisterModuleStep; #[async_trait] impl StepExecutor for RegisterModuleStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config_request: Arc = context.get_or_err("wasm_module_config")?; + let config_request: Arc = + context.get_or_err("wasm_module_config")?; let app_context: Arc = context.get_or_err("app_context")?; let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?; let file_size_bytes: Arc = context.get_or_err("file_size_bytes")?; diff --git a/sgl-model-gateway/src/core/workflow/steps/worker_registration.rs b/sgl-model-gateway/src/core/workflow/steps/worker_registration.rs index 2fcb52f48..b7bacbb20 100644 --- a/sgl-model-gateway/src/core/workflow/steps/worker_registration.rs +++ b/sgl-model-gateway/src/core/workflow/steps/worker_registration.rs @@ -464,7 +464,8 @@ impl StepExecutor for CreateWorkerStep { let config: Arc = context.get_or_err("worker_config")?; let app_context: Arc = context.get_or_err("app_context")?; let connection_mode: Arc = context.get_or_err("connection_mode")?; - let discovered_labels: Arc> = context.get_or_err("discovered_labels")?; + let discovered_labels: Arc> = + context.get_or_err("discovered_labels")?; // Check if worker already exists if app_context diff --git a/sgl-model-gateway/src/core/workflow/steps/worker_removal.rs b/sgl-model-gateway/src/core/workflow/steps/worker_removal.rs index 0f279e416..59e9607cf 100644 --- a/sgl-model-gateway/src/core/workflow/steps/worker_removal.rs +++ b/sgl-model-gateway/src/core/workflow/steps/worker_removal.rs @@ -106,7 +106,8 @@ pub struct RemoveFromPolicyRegistryStep; impl StepExecutor for RemoveFromPolicyRegistryStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { let app_context: Arc = context.get_or_err("app_context")?; - let workers_to_remove: Arc>> = context.get_or_err("workers_to_remove")?; + let workers_to_remove: Arc>> = + context.get_or_err("workers_to_remove")?; debug!( "Removing {} worker(s) from policy registry", diff --git a/sgl-model-gateway/src/core/workflow/types.rs b/sgl-model-gateway/src/core/workflow/types.rs index f9214f09b..90739d00a 100644 --- a/sgl-model-gateway/src/core/workflow/types.rs +++ b/sgl-model-gateway/src/core/workflow/types.rs @@ -220,7 +220,7 @@ impl WorkflowContext { /// Retrieve a value from the context, returning an error if not found pub fn get_or_err(&self, key: &str) -> Result, WorkflowError> { - self.get(key) + self.get::(key) .ok_or_else(|| WorkflowError::ContextValueNotFound(key.to_string())) } diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index ef721eb3c..82c312eb3 100644 --- a/sgl-model-gateway/src/routers/http/router.rs +++ b/sgl-model-gateway/src/routers/http/router.rs @@ -566,7 +566,6 @@ impl Router { // For streaming with load tracking, we need to manually decrement when done // Clone the worker Arc for the async block instead of looking it up again let stream_worker = worker.clone(); - let worker_url_owned = worker_url.to_string(); // Preserve headers for streaming response let mut response_headers = header_utils::preserve_response_headers(res.headers());