fix(worker): honor explicit worker connection schemes#1485
Open
heymrbox wants to merge 2 commits into
Open
Conversation
📝 WalkthroughWalkthroughThis PR extends gRPC over TLS ( ChangesgRPC over TLS Support
Sequence DiagramsequenceDiagram
participant Client as Caller
participant Detect as DetectConnectionModeStep
participant Util as util helpers
participant HTTP as HTTP health check
participant GRPC as gRPC health check
Client->>Detect: execute(url)
Detect->>Util: explicit_connection_mode(url)
alt explicit HTTP
Detect->>HTTP: http_health_url -> probe /health
HTTP-->>Detect: success/fail
else explicit gRPC (grpc/grpcs)
Detect->>GRPC: grpc_reachable_url -> gRPC health checks
GRPC-->>Detect: success/fail
else none (bare host:port)
Detect->>HTTP: probe
Detect->>GRPC: probe
end
Detect-->>Client: set connection_mode or error
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces explicit URL scheme detection in the connection detection workflow, allowing the system to honor grpc://, http://, and https:// prefixes. When a scheme is present, the system performs a targeted health check instead of probing both protocols in parallel. Feedback suggests replacing the manual string prefix matching with a proper URL parsing library to improve reliability, handle case sensitivity more robustly, and avoid potential issues with relative URLs.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Clean, well-structured change. The explicit scheme detection is a good UX improvement over always probing both protocols. Tests cover the key cases.
Two minor nits posted inline (both 🟡):
- Case normalization mismatch:
to_ascii_lowercase()in the detector accepts mixed-case schemes that downstream functions can't handle — suggest matching exact lowercase prefixes instead. - Redundant protocol variable: error message includes both
{protocol}and{connection_mode}Display — one suffices.
Member
CatherineSue
left a comment
CatherineSue
left a comment
There was a problem hiding this comment.
The fix LGTM.
I think it would be also better to make try_http_reachable and try_grpc_reachable more robust. The bug exists because they both try to use strip_protocol. Maybe they both need to be more strict. Such as:
try_http_reachable:
accepts http://, https://, or bare host:port
rejects grpc://, grpcs://
try_grpc_reachable:
accepts grpc://, grpcs://, or bare host:port
rejects http://, https://
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
heymrbox
force-pushed
the
weiwei/fix-worker
branch
from
fddb01a to
8aabf5c
Compare
Signed-off-by: Weiwei Zheng <heymrbox@gmail.com>
heymrbox
force-pushed
the
weiwei/fix-worker
branch
from
8aabf5c to
c652b76
Compare
There was a problem hiding this comment.
Actionable comments posted: 9
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/grpc_client/src/mlx_engine.rs`:
- Around line 125-131: Add a new async Tokio test named
test_grpcs_scheme_conversion next to the existing
test_client_connect_invalid_endpoint to assert the new scheme conversion; call
MlxEngineClient::connect("grpcs://localhost:50051").await and assert the result
is Err (connection failure expected) so the failure is due to connection and not
scheme parsing, ensuring the grpcs:// → https:// path in the endpoint
normalization branch is covered.
- Around line 125-131: The scheme-prefix checks for building http_endpoint in
MlxEngineClient::connect() (and connect_with_trace_injector()) are
case-sensitive; make them case-insensitive by lowercasing a temporary copy of
endpoint (e.g., let lower = endpoint.to_lowercase()) and use
lower.strip_prefix("grpc://") / lower.strip_prefix("grpcs://") to detect the
scheme, but when constructing the replacement URL use the substring of the
original endpoint after the matched prefix length (so host/port/case are
preserved) to build "http://{rest}" or "https://{rest}"; apply this change in
the code that sets http_endpoint so user-supplied schemes like GRPC:// or
GrPcS:// are handled correctly.
In `@crates/grpc_client/src/sglang_scheduler.rs`:
- Around line 146-152: Add a unit test that mirrors
test_client_connect_invalid_endpoint() to cover the new grpcs:// scheme: call
the same code path that computes http_endpoint (the logic using
endpoint.strip_prefix in sglang_scheduler.rs) with a "grpcs://host:port" input
and assert the resulting http_endpoint is "https://host:port"; ensure the test
lives in the same test module as test_client_connect_invalid_endpoint() so it
exercises the same conversion logic and fails if the grpcs:// -> https://
mapping regresses.
- Around line 146-152: The scheme check in sglang_scheduler.rs uses
case-sensitive strip_prefix on endpoint, so endpoints like "GRPC://" or
"GRPCS://" are not recognized; update the logic around http_endpoint to do
case-insensitive matching by comparing a lowercase copy (e.g., let lower =
endpoint.to_ascii_lowercase()) and then if lower.strip_prefix("grpc://") /
"grpcs://", derive the addr by slicing the original endpoint past the matched
prefix (so original casing of the host is preserved) and format into
"http://{addr}" or "https://{addr}" accordingly; update the block that sets
http_endpoint to use this case-insensitive check on endpoint while still using
the original endpoint string for the formatted address.
In `@crates/grpc_client/src/trtllm_service.rs`:
- Around line 145-151: Add a unit test that exercises the scheme conversion
logic around endpoint.strip_prefix so grpcs:// URIs are converted to https://;
specifically, create a test that supplies an endpoint like
"grpcs://example.com:1234" and asserts the resulting http_endpoint equals
"https://example.com:1234" (similar to existing tests for grpc:// → http://).
Place the test adjacent to other tests for the module and reference the same
conversion code that computes http_endpoint (the block using
endpoint.strip_prefix and format!).
- Around line 145-151: The scheme matching for endpoint is case-sensitive and
misses mixed-case schemes like "GRPCS://"; update the logic in the connect flow
(where endpoint and http_endpoint are computed, e.g. GrpcClient::connect in
trtllm_service.rs) to perform case-insensitive matching by using
endpoint.to_ascii_lowercase() for the prefix checks and then slicing the
original endpoint to get the addr (so the host/port retains original casing).
Concretely: compute a lower = endpoint.to_ascii_lowercase(), check
lower.strip_prefix("grpc://") / lower.strip_prefix("grpcs://"), and when a match
is found use the corresponding byte/char index to take the addr from the
original endpoint and format "http://{addr}" or "https://{addr}" into
http_endpoint.
In `@crates/grpc_client/src/vllm_engine.rs`:
- Around line 146-152: The new scheme conversion in vllm_engine.rs (the endpoint
normalization that switches "grpc://"→"http://" and "grpcs://"→"https://" inside
the block that builds http_endpoint) lacks a unit test; add a test (e.g.,
test_endpoint_scheme_conversion or test_grpcs_to_https_conversion) that calls
the function or constructs the same normalization path with an input
"grpcs://example.com:1234" and asserts the resulting http_endpoint equals
"https://example.com:1234"; place the test near other vllm_engine tests,
exercise both "grpc://" and "grpcs://" cases, and run cargo test to ensure
coverage.
- Around line 146-152: The scheme check for building http_endpoint is
case-sensitive and misses mixed-case prefixes; fix by computing a lowercase copy
(e.g., let lower = endpoint.to_ascii_lowercase()) and test strip_prefix against
"grpc://" and "grpcs://" on that lowercase string, then when a match is found
derive the addr by slicing the original endpoint string past the matched prefix
length (so you preserve the original rest of the URL) and format to
"http://{addr}" or "https://{addr}"; update the logic around http_endpoint and
endpoint.strip_prefix to use this lowercased check and original slicing.
In `@model_gateway/src/workflow/steps/local/detect_connection.rs`:
- Around line 19-27: The explicit_connection_mode function currently uses
case-sensitive starts_with checks and should detect schemes case-insensitively
to match util.rs behavior; update explicit_connection_mode to compare the URL
scheme using eq_ignore_ascii_case (or lowercase the prefix) for "grpc", "grpcs",
"http", and "https" so mixed-case schemes like "GRPC://..." are recognized as
ConnectionMode::Grpc/Http, and add a unit test for a mixed-case URL (e.g.,
"GRPC://localhost:30001") to cover this case.
🪄 Autofix (Beta)
✅ Autofix completed
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 3043b542-5812-4560-8b71-b1c1d1c00b40
📒 Files selected for processing (6)
crates/grpc_client/src/mlx_engine.rscrates/grpc_client/src/sglang_scheduler.rscrates/grpc_client/src/trtllm_service.rscrates/grpc_client/src/vllm_engine.rsmodel_gateway/src/workflow/steps/local/detect_connection.rsmodel_gateway/src/workflow/steps/util.rs
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c652b76f73
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Comment on lines
+20
to
+21
| if url.starts_with("grpc://") || url.starts_with("grpcs://") { | ||
| Some(ConnectionMode::Grpc) |
There was a problem hiding this comment.
Treating grpcs://... as an explicit gRPC URL here makes the registration workflow proceed, but downstream URL normalization still only whitelists http://, https://, and grpc:// (create_worker.rs normalize_url), so a successful grpcs://host:port registration is rewritten to grpc://grpcs://host:port. That malformed URL is then used by GrpcClient::connect, causing runtime connection failures for the newly registered worker.
Useful? React with 👍 / 👎.
|
Note Autofix is a beta feature. Expect some limitations and changes as we gather feedback and continue to improve it. Fixes Applied SuccessfullyFixed 4 file(s) based on 4 unresolved review comments. Files modified:
Commit: The changes have been pushed to the Time taken: |
Fixed 4 file(s) based on 4 unresolved review comments. Co-authored-by: CodeRabbit <noreply@coderabbit.ai>
There was a problem hiding this comment.
♻️ Duplicate comments (4)
crates/grpc_client/src/mlx_engine.rs (1)
125-131:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winImplement case-insensitive scheme matching.
The scheme checks use case-sensitive
strip_prefixand won't handle mixed-case schemes likeGRPC://orGrPcS://. Past review analysis showsworker.rs:952passesmetadata.spec.urldirectly toGrpcClient::connect(), and sincegrpc_base_urlinutil.rspreserves the original URL casing when a scheme is detected, uppercase schemes will bypass the conversion and cause connection failures.🔧 Proposed fix
- let http_endpoint = if let Some(addr) = endpoint.strip_prefix("grpc://") { + let lower = endpoint.to_ascii_lowercase(); + let http_endpoint = if let Some(_) = lower.strip_prefix("grpc://") { + let addr = &endpoint[7..]; format!("http://{addr}") - } else if let Some(addr) = endpoint.strip_prefix("grpcs://") { + } else if let Some(_) = lower.strip_prefix("grpcs://") { + let addr = &endpoint[8..]; format!("https://{addr}") } else { endpoint.to_string() };🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/grpc_client/src/mlx_engine.rs` around lines 125 - 131, The scheme-matching for building http_endpoint in the mlx_engine.rs block (the if/else that currently uses endpoint.strip_prefix("grpc://") and strip_prefix("grpcs://")) is case-sensitive and fails for mixed/upper-case schemes; update this block to perform case-insensitive detection (e.g. compare a lowercased view of endpoint or use a case-insensitive starts-with) and then strip only the scheme portion while preserving the original remainder of endpoint so GrpcClient::connect() (and grpc_base_url in util.rs) receive the correct host/path; adjust the logic around the http_endpoint variable and the endpoint handling to use these case-insensitive checks.crates/grpc_client/src/vllm_engine.rs (1)
145-152:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winImplement case-insensitive scheme matching.
The scheme checks use case-sensitive
strip_prefixand won't handle mixed-case schemes likeGRPC://orGrPcS://. Past review analysis shows that direct callers bypass workflow layer normalization, and sincegrpc_base_urlinutil.rspreserves the original URL casing when a scheme is detected, mixed-case schemes will fail the conversion and cause connection failures.🔧 Proposed fix
// Convert gRPC schemes to tonic-compatible HTTP(S) endpoints. - let http_endpoint = if let Some(addr) = endpoint.strip_prefix("grpc://") { + let lower = endpoint.to_ascii_lowercase(); + let http_endpoint = if let Some(_) = lower.strip_prefix("grpc://") { + let addr = &endpoint[7..]; format!("http://{addr}") - } else if let Some(addr) = endpoint.strip_prefix("grpcs://") { + } else if let Some(_) = lower.strip_prefix("grpcs://") { + let addr = &endpoint[8..]; format!("https://{addr}") } else { endpoint.to_string() };🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/grpc_client/src/vllm_engine.rs` around lines 145 - 152, The current scheme matching on `endpoint` uses case-sensitive `strip_prefix` and fails for mixed-case schemes; update the conversion that sets `http_endpoint` so it matches schemes case-insensitively by comparing a lowercased view (e.g., `let lower = endpoint.to_lowercase()`) and, when a match is found, take the remainder from the original `endpoint` by slicing with the scheme length (so casing of the address is preserved) and then format to "http://{addr}" or "https://{addr}". Make this change in the block that computes `http_endpoint` (the code referencing `endpoint` and producing `http_endpoint`) so both "grpc://" and "grpcs://" are handled regardless of case.crates/grpc_client/src/sglang_scheduler.rs (1)
145-152:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winImplement case-insensitive scheme matching.
The scheme checks use case-sensitive
strip_prefixand won't handle mixed-case schemes likeGRPC://orGrPcS://. Past review analysis confirms thatworker.rs:952and other callers pass URLs directly without normalization, and mixed-case schemes will bypass the conversion logic causing connection failures.🔧 Proposed fix
// Convert gRPC schemes to tonic-compatible HTTP(S) endpoints. - let http_endpoint = if let Some(addr) = endpoint.strip_prefix("grpc://") { + let lower = endpoint.to_ascii_lowercase(); + let http_endpoint = if let Some(_) = lower.strip_prefix("grpc://") { + let addr = &endpoint[7..]; format!("http://{addr}") - } else if let Some(addr) = endpoint.strip_prefix("grpcs://") { + } else if let Some(_) = lower.strip_prefix("grpcs://") { + let addr = &endpoint[8..]; format!("https://{addr}") } else { endpoint.to_string() };🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/grpc_client/src/sglang_scheduler.rs` around lines 145 - 152, The scheme checks on endpoint use case-sensitive strip_prefix and miss mixed-case schemes; update the conversion that sets http_endpoint to do case-insensitive matching: create a lowercase copy (e.g., let lower = endpoint.to_ascii_lowercase()), test lower.strip_prefix("grpc://") and lower.strip_prefix("grpcs://"), and when a match is found compute the original address suffix by slicing the original endpoint past the matched scheme length (so you preserve original casing/host) and then format!("http://{addr}") or format!("https://{addr}") accordingly; update the logic that sets http_endpoint (same variable) and reuse the original endpoint variable for the suffix.crates/grpc_client/src/trtllm_service.rs (1)
144-151:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winImplement case-insensitive scheme matching.
The scheme checks use case-sensitive
strip_prefixand won't handle mixed-case schemes likeGRPC://orGrPcS://. Past review analysis confirms that direct callers likeworker.rs:952pass URLs without normalization, and sincegrpc_base_urlpreserves original URL casing, mixed-case schemes will fail the conversion and cause connection errors.🔧 Proposed fix
// Convert gRPC schemes to tonic-compatible HTTP(S) endpoints. - let http_endpoint = if let Some(addr) = endpoint.strip_prefix("grpc://") { + let lower = endpoint.to_ascii_lowercase(); + let http_endpoint = if let Some(_) = lower.strip_prefix("grpc://") { + let addr = &endpoint[7..]; format!("http://{addr}") - } else if let Some(addr) = endpoint.strip_prefix("grpcs://") { + } else if let Some(_) = lower.strip_prefix("grpcs://") { + let addr = &endpoint[8..]; format!("https://{addr}") } else { endpoint.to_string() };🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/grpc_client/src/trtllm_service.rs` around lines 144 - 151, The current scheme matching uses case-sensitive strip_prefix on endpoint and will miss mixed-case schemes; fix by computing a lowercase copy (e.g. let lower = endpoint.to_ascii_lowercase()) and perform strip_prefix checks against lower (e.g. lower.strip_prefix("grpc://") / lower.strip_prefix("grpcs://")), then map the returned suffix back to the original endpoint string to build http_endpoint so original host casing is preserved; update the logic that sets http_endpoint (using endpoint and the strip_prefix results) to use these case-insensitive checks instead of direct strip_prefix calls.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Duplicate comments:
In `@crates/grpc_client/src/mlx_engine.rs`:
- Around line 125-131: The scheme-matching for building http_endpoint in the
mlx_engine.rs block (the if/else that currently uses
endpoint.strip_prefix("grpc://") and strip_prefix("grpcs://")) is case-sensitive
and fails for mixed/upper-case schemes; update this block to perform
case-insensitive detection (e.g. compare a lowercased view of endpoint or use a
case-insensitive starts-with) and then strip only the scheme portion while
preserving the original remainder of endpoint so GrpcClient::connect() (and
grpc_base_url in util.rs) receive the correct host/path; adjust the logic around
the http_endpoint variable and the endpoint handling to use these
case-insensitive checks.
In `@crates/grpc_client/src/sglang_scheduler.rs`:
- Around line 145-152: The scheme checks on endpoint use case-sensitive
strip_prefix and miss mixed-case schemes; update the conversion that sets
http_endpoint to do case-insensitive matching: create a lowercase copy (e.g.,
let lower = endpoint.to_ascii_lowercase()), test lower.strip_prefix("grpc://")
and lower.strip_prefix("grpcs://"), and when a match is found compute the
original address suffix by slicing the original endpoint past the matched scheme
length (so you preserve original casing/host) and then format!("http://{addr}")
or format!("https://{addr}") accordingly; update the logic that sets
http_endpoint (same variable) and reuse the original endpoint variable for the
suffix.
In `@crates/grpc_client/src/trtllm_service.rs`:
- Around line 144-151: The current scheme matching uses case-sensitive
strip_prefix on endpoint and will miss mixed-case schemes; fix by computing a
lowercase copy (e.g. let lower = endpoint.to_ascii_lowercase()) and perform
strip_prefix checks against lower (e.g. lower.strip_prefix("grpc://") /
lower.strip_prefix("grpcs://")), then map the returned suffix back to the
original endpoint string to build http_endpoint so original host casing is
preserved; update the logic that sets http_endpoint (using endpoint and the
strip_prefix results) to use these case-insensitive checks instead of direct
strip_prefix calls.
In `@crates/grpc_client/src/vllm_engine.rs`:
- Around line 145-152: The current scheme matching on `endpoint` uses
case-sensitive `strip_prefix` and fails for mixed-case schemes; update the
conversion that sets `http_endpoint` so it matches schemes case-insensitively by
comparing a lowercased view (e.g., `let lower = endpoint.to_lowercase()`) and,
when a match is found, take the remainder from the original `endpoint` by
slicing with the scheme length (so casing of the address is preserved) and then
format to "http://{addr}" or "https://{addr}". Make this change in the block
that computes `http_endpoint` (the code referencing `endpoint` and producing
`http_endpoint`) so both "grpc://" and "grpcs://" are handled regardless of
case.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 8bec4bca-a9c8-456a-9ee7-5f3ca1680480
📒 Files selected for processing (6)
crates/grpc_client/src/mlx_engine.rscrates/grpc_client/src/sglang_scheduler.rscrates/grpc_client/src/trtllm_service.rscrates/grpc_client/src/vllm_engine.rsmodel_gateway/src/workflow/steps/local/detect_connection.rsmodel_gateway/src/workflow/steps/util.rs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Problem
Explicit worker URL schemes were not fully respected during local connection mode detection. A worker configured with
grpc://host:portcould still be probed over HTTP, and if the HTTP health check succeeded, the worker could be classified asConnectionMode::Httpdespite the explicit gRPC scheme.Solution
Honor explicit URL schemes before falling back to auto-detection:
grpc://...runs only the gRPC reachability check and setsConnectionMode::Grpchttp://...andhttps://...run only the HTTP reachability check and setConnectionMode::Httphost:portURLs keep the existing behavior of probing both protocols, with HTTP priority if both passChanges
Test Plan
Reproduced the issue by reviewing
DetectConnectionModeStep: before this change, all URLs were passed into both HTTP and gRPC probes, sogrpc://...could still be selected as HTTP when/healthresponded.After this change, explicit schemes are handled before the auto-detection block, so
grpc://...no longer falls through to HTTP probing.Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit