[model-gateway] fix annotation error and code formating (#14910)

This commit is contained in:
Simo Lin
2025-12-11 07:06:32 -08:00
committed by GitHub
parent 5b0872d2d7
commit cd1c1fa5bc
6 changed files with 25 additions and 14 deletions
@@ -46,7 +46,8 @@ pub struct ConnectMcpServerStep;
#[async_trait] #[async_trait]
impl StepExecutor for ConnectMcpServerStep { impl StepExecutor for ConnectMcpServerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?; let config_request: Arc<McpServerConfigRequest> =
context.get_or_err("mcp_server_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?; let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
debug!("Connecting to MCP server: {}", config_request.name); debug!("Connecting to MCP server: {}", config_request.name);
@@ -98,7 +99,8 @@ impl StepExecutor for DiscoverMcpInventoryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
use rmcp::{service::RunningService, RoleClient}; use rmcp::{service::RunningService, RoleClient};
let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?; let config_request: Arc<McpServerConfigRequest> =
context.get_or_err("mcp_server_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?; let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let mcp_client: Arc<RunningService<RoleClient, ()>> = context.get_or_err("mcp_client")?; let mcp_client: Arc<RunningService<RoleClient, ()>> = context.get_or_err("mcp_client")?;
@@ -143,7 +145,8 @@ impl StepExecutor for RegisterMcpServerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
use rmcp::{service::RunningService, RoleClient}; use rmcp::{service::RunningService, RoleClient};
let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?; let config_request: Arc<McpServerConfigRequest> =
context.get_or_err("mcp_server_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?; let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let mcp_client: Arc<RunningService<RoleClient, ()>> = context.get_or_err("mcp_client")?; let mcp_client: Arc<RunningService<RoleClient, ()>> = context.get_or_err("mcp_client")?;
@@ -183,7 +186,8 @@ pub struct ValidateRegistrationStep;
#[async_trait] #[async_trait]
impl StepExecutor for ValidateRegistrationStep { impl StepExecutor for ValidateRegistrationStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<McpServerConfigRequest> = context.get_or_err("mcp_server_config")?; let config_request: Arc<McpServerConfigRequest> =
context.get_or_err("mcp_server_config")?;
let client_registered = context let client_registered = context
.get::<RunningService<RoleClient, ()>>("mcp_client") .get::<RunningService<RoleClient, ()>>("mcp_client")
@@ -75,7 +75,8 @@ pub struct ValidateDescriptorStep;
#[async_trait] #[async_trait]
impl StepExecutor for ValidateDescriptorStep { impl StepExecutor for ValidateDescriptorStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?; let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let descriptor = &config_request.descriptor; let descriptor = &config_request.descriptor;
@@ -233,7 +234,8 @@ pub struct CalculateHashStep;
#[async_trait] #[async_trait]
impl StepExecutor for CalculateHashStep { impl StepExecutor for CalculateHashStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?; let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let file_path = &config_request.descriptor.file_path; let file_path = &config_request.descriptor.file_path;
@@ -291,7 +293,8 @@ pub struct CheckDuplicateStep;
#[async_trait] #[async_trait]
impl StepExecutor for CheckDuplicateStep { impl StepExecutor for CheckDuplicateStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?; let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?; let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?; let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?;
@@ -339,7 +342,8 @@ pub struct LoadWasmBytesStep;
#[async_trait] #[async_trait]
impl StepExecutor for LoadWasmBytesStep { impl StepExecutor for LoadWasmBytesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?; let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let file_path = &config_request.descriptor.file_path; let file_path = &config_request.descriptor.file_path;
@@ -374,7 +378,8 @@ pub struct ValidateWasmComponentStep;
#[async_trait] #[async_trait]
impl StepExecutor for ValidateWasmComponentStep { impl StepExecutor for ValidateWasmComponentStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?; let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let wasm_bytes: Arc<Vec<u8>> = context.get_or_err("wasm_bytes")?; let wasm_bytes: Arc<Vec<u8>> = context.get_or_err("wasm_bytes")?;
debug!( debug!(
@@ -425,7 +430,8 @@ pub struct RegisterModuleStep;
#[async_trait] #[async_trait]
impl StepExecutor for RegisterModuleStep { impl StepExecutor for RegisterModuleStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> = context.get_or_err("wasm_module_config")?; let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?; let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?; let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?;
let file_size_bytes: Arc<u64> = context.get_or_err("file_size_bytes")?; let file_size_bytes: Arc<u64> = context.get_or_err("file_size_bytes")?;
@@ -464,7 +464,8 @@ impl StepExecutor for CreateWorkerStep {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?; let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?; let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?; let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
let discovered_labels: Arc<HashMap<String, String>> = context.get_or_err("discovered_labels")?; let discovered_labels: Arc<HashMap<String, String>> =
context.get_or_err("discovered_labels")?;
// Check if worker already exists // Check if worker already exists
if app_context if app_context
@@ -106,7 +106,8 @@ pub struct RemoveFromPolicyRegistryStep;
impl StepExecutor for RemoveFromPolicyRegistryStep { impl StepExecutor for RemoveFromPolicyRegistryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?; let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers_to_remove: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers_to_remove")?; let workers_to_remove: Arc<Vec<Arc<dyn Worker>>> =
context.get_or_err("workers_to_remove")?;
debug!( debug!(
"Removing {} worker(s) from policy registry", "Removing {} worker(s) from policy registry",
+1 -1
View File
@@ -220,7 +220,7 @@ impl WorkflowContext {
/// Retrieve a value from the context, returning an error if not found /// 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> { pub fn get_or_err<T: Send + Sync + 'static>(&self, key: &str) -> Result<Arc<T>, WorkflowError> {
self.get(key) self.get::<T>(key)
.ok_or_else(|| WorkflowError::ContextValueNotFound(key.to_string())) .ok_or_else(|| WorkflowError::ContextValueNotFound(key.to_string()))
} }
@@ -566,7 +566,6 @@ impl Router {
// For streaming with load tracking, we need to manually decrement when done // 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 // Clone the worker Arc for the async block instead of looking it up again
let stream_worker = worker.clone(); let stream_worker = worker.clone();
let worker_url_owned = worker_url.to_string();
// Preserve headers for streaming response // Preserve headers for streaming response
let mut response_headers = header_utils::preserve_response_headers(res.headers()); let mut response_headers = header_utils::preserve_response_headers(res.headers());