diff --git a/model_gateway/src/workflow/steps/local/detect_connection.rs b/model_gateway/src/workflow/steps/local/detect_connection.rs index aa984dd251..deffe5fd45 100644 --- a/model_gateway/src/workflow/steps/local/detect_connection.rs +++ b/model_gateway/src/workflow/steps/local/detect_connection.rs @@ -16,9 +16,21 @@ use crate::{ }, }; +fn explicit_connection_mode(url: &str) -> Option { + let normalized_url = url.to_ascii_lowercase(); + if normalized_url.starts_with("grpc://") { + Some(ConnectionMode::Grpc) + } else if normalized_url.starts_with("http://") || normalized_url.starts_with("https://") { + Some(ConnectionMode::Http) + } else { + None + } +} + /// Step 1: Detect connection mode (HTTP vs gRPC). /// -/// Probes both protocols in parallel. HTTP takes priority if both succeed. +/// Explicit URL schemes are honored. For bare host:port URLs, probes both +/// protocols in parallel and HTTP takes priority if both succeed. /// Does NOT detect backend runtime — that's handled by DetectBackendStep. pub struct DetectConnectionModeStep; @@ -51,6 +63,37 @@ impl StepExecutor for DetectConnectionModeStep { .unwrap_or(app_context.router_config.health_check.timeout_secs); let client = &app_context.client; + if let Some(connection_mode) = explicit_connection_mode(&url) { + let result = match connection_mode { + ConnectionMode::Http => try_http_reachable(&url, timeout, client).await, + ConnectionMode::Grpc => try_grpc_reachable(&url, timeout).await, + }; + + match result { + Ok(()) => { + debug!( + "{} explicitly configured as {}", + config.url, connection_mode + ); + context.data.connection_mode = Some(connection_mode); + return Ok(StepResult::Success); + } + Err(err) => { + let protocol = match connection_mode { + ConnectionMode::Http => "HTTP", + ConnectionMode::Grpc => "gRPC", + }; + return Err(WorkflowError::StepFailed { + step_id: StepId::new("detect_connection_mode"), + message: format!( + "{protocol} health check failed for explicit {} worker URL {}: {}", + connection_mode, config.url, err + ), + }); + } + } + } + let (http_result, grpc_result) = tokio::join!( try_http_reachable(&url, timeout, client), try_grpc_reachable(&url, timeout) @@ -84,3 +127,33 @@ impl StepExecutor for DetectConnectionModeStep { true } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn explicit_connection_mode_honors_grpc_scheme() { + assert_eq!( + explicit_connection_mode("grpc://localhost:30001"), + Some(ConnectionMode::Grpc) + ); + } + + #[test] + fn explicit_connection_mode_honors_http_schemes() { + assert_eq!( + explicit_connection_mode("http://localhost:30000"), + Some(ConnectionMode::Http) + ); + assert_eq!( + explicit_connection_mode("https://example.com"), + Some(ConnectionMode::Http) + ); + } + + #[test] + fn explicit_connection_mode_leaves_bare_urls_for_probe_detection() { + assert_eq!(explicit_connection_mode("localhost:30000"), None); + } +}