Skip to content
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 74 additions & 1 deletion model_gateway/src/workflow/steps/local/detect_connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,21 @@ use crate::{
},
};

fn explicit_connection_mode(url: &str) -> Option<ConnectionMode> {
let normalized_url = url.to_ascii_lowercase();

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: to_ascii_lowercase() makes detection case-insensitive (e.g. GRPC://host is recognized), but the original mixed-case URL is then passed unchanged to try_grpc_reachable / try_http_reachable, which use case-sensitive starts_with("grpc://") / starts_with("https://") internally. For a URL like GRPC://localhost:30001, detection correctly picks gRPC, but try_grpc_reachable falls into its else branch and produces a malformed grpc://GRPC://localhost:30001.

Two options:

  1. Drop the to_ascii_lowercase() so detection is case-sensitive like downstream (simplest, mixed-case schemes are rare).
  2. Normalize the URL itself (lowercase the scheme portion) and pass the normalized URL to the probe functions.

Either way, the detection and the downstream consumers should agree on case handling.

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;

Expand Down Expand Up @@ -51,6 +63,37 @@ impl StepExecutor<WorkerWorkflowData> 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,
};
Comment on lines +66 to +70

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The current implementation is susceptible to case-sensitivity issues because the underlying reachability probes perform case-sensitive prefix matching on the URL. To reliably handle schemes and distinguish between different URL types, parse the URL using a dedicated library and check its scheme instead of using string prefix matching. This approach avoids misclassifying URLs and ensures robustness against variations like 'GRPC://'.

References
  1. To reliably distinguish data URIs from other URLs, parse the URL and check its scheme instead of using string prefix matching. This avoids misclassifying relative URLs.


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
),
});
Comment on lines +82 to +92

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The error message is redundant, repeating "health check failed" and including both the capitalized protocol string and the connection_mode enum. Additionally, the underlying error err from the probe already includes a "Health check failed" prefix. A more concise message would improve log clarity.

Suggested change
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 protocol = match connection_mode {
ConnectionMode::Http => "HTTP",
ConnectionMode::Grpc => "gRPC",
};
return Err(WorkflowError::StepFailed {
step_id: StepId::new("detect_connection_mode"),
message: format!(
"Explicit {protocol} worker URL {} is unreachable: {err}",
config.url
),
});

}
}
}

let (http_result, grpc_result) = tokio::join!(
try_http_reachable(&url, timeout, client),
try_grpc_reachable(&url, timeout)
Expand Down Expand Up @@ -84,3 +127,33 @@ impl StepExecutor<WorkerWorkflowData> 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);
}
}
Loading