fix(gateway): use URL-based detection to eliminate classify race condition - #839
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the robustness of worker classification within the gateway by addressing a critical race condition. The changes introduce a more deterministic and reliable approach to distinguish between local and external worker types, primarily through URL-based detection and explicit Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughWorker classification now respects explicit Changes
Sequence Diagram(s)sequenceDiagram
participant Gateway as Gateway
participant Worker as Worker (endpoint)
participant Heuristics as URL/Runtime Heuristics
participant Prober as Prober (/v1/models, /health, gRPC)
Gateway->>Worker: receive add_worker(config)
Gateway->>Heuristics: is runtime_type specified?
alt specified && runtime_type == External
Heuristics-->>Gateway: classify External
else specified && runtime_type != External
Heuristics-->>Gateway: classify Local
else unspecified
Heuristics->>Heuristics: ProviderType::from_url(config.url)
alt URL matches known cloud provider
Heuristics-->>Gateway: classify External
else
Gateway->>Prober: GET /v1/models -> parse owned_by
Prober-->>Gateway: owned_by or None
alt owned_by in {sglang, vllm, trtllm}
Gateway-->>Worker: classify Local
else
Gateway->>Prober: probe /health + gRPC health
alt health OK
Prober-->>Gateway: health OK -> classify Local
else
Gateway-->>Worker: classify Local (fallback)
end
end
end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
…ition The classify step previously treated a /v1/models response (without /health) as proof of an External worker. This caused a race condition: local backends (sglang, vllm) that serve /v1/models before /health is ready were misclassified as External. What changed: - classify.rs: replace /v1/models reachability check with two safer signals: (a) URL-based detection via ProviderType::from_url() for known cloud providers (openai.com, anthropic.com, x.ai, googleapis.com) and (b) /v1/models owned_by field matching known local backends (sglang, vllm, trtllm). Combine the explicit External and explicit Local branches into a single is_specified() check. New detection logic: 1. Any explicit runtime → classify immediately (External or Local) 2. URL matches known cloud provider → External 3. /health responds → Local 4. gRPC health responds → Local 5. /v1/models owned_by matches local backend → Local 6. Nothing conclusive → default Local Why: the old heuristic (/v1/models responds + no /health = External) was unreliable during backend startup. URL-based detection is instant, deterministic, and has zero race conditions. For external providers on private IPs (e.g., a proxy to OpenAI), users must set runtime: external explicitly — this is a reasonable requirement since URL detection cannot identify them. Refs: #820 Signed-off-by: Simon Lin <simonslin@gmail.com> Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
20a1f62 to
bb5af00
Compare
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request effectively addresses a race condition in worker classification by introducing a more robust detection logic. The new approach prioritizes explicit configuration, uses URL-based heuristics for known external providers, and inspects the /v1/models response to identify local backends, which is a significant improvement over the previous, less reliable method. The changes are well-documented and the trade-offs are clearly explained. I have one suggestion to further improve the code's robustness and maintainability.
| let body: serde_json::Value = resp.json().await.ok()?; | ||
| let owned_by = body | ||
| .get("data")? | ||
| .as_array()? | ||
| .first()? | ||
| .get("owned_by")? | ||
| .as_str()? | ||
| .to_lowercase(); |
There was a problem hiding this comment.
The current implementation uses serde_json::Value and a chain of ? operators to extract the owned_by field. While this works, it can be brittle and harder to read. Using dedicated structs for deserialization would be more robust, type-safe, and self-documenting, clearly defining the expected structure of the JSON response.
#[derive(serde::Deserialize)]
struct Model {
owned_by: String,
}
#[derive(serde::Deserialize)]
struct ModelList {
data: Vec<Model>,
}
let models: ModelList = resp.json().await.ok()?;
let owned_by = models.data.first()?.owned_by.to_lowercase();![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: bb5af00bd1
ℹ️ 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".
| // Explicit local runtime (Sglang, Vllm, Trtllm) → Local | ||
| if config.runtime_type.is_specified() { | ||
| // 3. URL matches known cloud provider → External (no probing needed) | ||
| if let Some(provider) = ProviderType::from_url(&config.url) { |
There was a problem hiding this comment.
Classifying as External purely on ProviderType::from_url(&config.url) here can produce false positives because from_url currently matches hosts with ends_with("openai.com")/similar, so domains like notopenai.com are treated as OpenAI. In this step that means we return early as External and skip /health/gRPC probes, so a local worker on such a hostname is misclassified and routed through the wrong workflow. Please add a strict host/domain-boundary check before using URL-based auto-classification in this branch.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
model_gateway/src/core/steps/worker/classify.rs (1)
140-158: 🧹 Nitpick | 🔵 Trivial
/v1/modelsis now only influencing the debug path.Lines 141-149 and Lines 152-158 both end in
WorkerKind::Local, so this request cannot change the classification anymore. On the slow path it just adds another timeout window and authenticated/v1/modelscall during registration. Either drop the probe or use it to differentiate a real non-local case.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/steps/worker/classify.rs` around lines 140 - 158, The /v1/models probe (probe_models_owned_by) currently only ever leads to setting context.data.worker_kind = Some(WorkerKind::Local), duplicating the fallback and only adding extra latency; remove the probe_models_owned_by call or make it meaningful by using its result to choose a non-local kind: call probe_models_owned_by(&config.url, ...) and if it returns Some(owned_by) then set context.data.worker_kind = Some(WorkerKind::Local) only when owned_by matches the local identifier (otherwise set context.data.worker_kind = Some(WorkerKind::Remote) or another appropriate non-local enum variant), and ensure you return Ok(StepResult::Success) after setting the kind so the slow-path timeout is not incurred unnecessarily.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/core/steps/worker/classify.rs`:
- Around line 105-111: ProviderType::from_url currently uses host.ends_with(...)
which false-positives lookalike domains; change its matching logic in
crates/protocols/src/worker.rs (the ProviderType::from_url helper) to require
exact host equality or a dot-prefixed suffix (e.g., host == "openai.com" ||
host.ends_with(".openai.com")) for each provider, and update all provider checks
(OpenAI, x.ai, googleapis/generative language) accordingly so subdomains match
but lookalikes like myopenai.com or notgoogleapis.com do not. Also add
regression tests covering the specified boundary cases (api.openai.com → match,
myopenai.com → no match, api.x.ai → match, teamx.ai → no match,
generativelanguage.googleapis.com → match, notgoogleapis.com → no match) to
validate the fix.
---
Outside diff comments:
In `@model_gateway/src/core/steps/worker/classify.rs`:
- Around line 140-158: The /v1/models probe (probe_models_owned_by) currently
only ever leads to setting context.data.worker_kind = Some(WorkerKind::Local),
duplicating the fallback and only adding extra latency; remove the
probe_models_owned_by call or make it meaningful by using its result to choose a
non-local kind: call probe_models_owned_by(&config.url, ...) and if it returns
Some(owned_by) then set context.data.worker_kind = Some(WorkerKind::Local) only
when owned_by matches the local identifier (otherwise set
context.data.worker_kind = Some(WorkerKind::Remote) or another appropriate
non-local enum variant), and ensure you return Ok(StepResult::Success) after
setting the kind so the slow-path timeout is not incurred unnecessarily.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 57a64d90-2494-4bdf-be0e-8764a5de36ac
📒 Files selected for processing (1)
model_gateway/src/core/steps/worker/classify.rs
| if let Some(provider) = ProviderType::from_url(&config.url) { | ||
| debug!( | ||
| "Worker {} explicitly configured as {} → Local", | ||
| config.url, config.runtime_type | ||
| "Worker {} URL matches known provider ({}) → External", | ||
| config.url, provider | ||
| ); | ||
| context.data.worker_kind = Some(WorkerKind::Local); | ||
| context.data.worker_kind = Some(WorkerKind::External); | ||
| return Ok(StepResult::Success); |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== ProviderType::from_url implementation =="
fd -0 '^worker.rs$' | xargs -0 rg -n -C2 \
'pub fn from_url|ends_with\("openai\.com"\)|ends_with\("x\.ai"\)|ends_with\("anthropic\.com"\)|ends_with\("googleapis\.com"\)'
echo
echo "== Boundary examples for raw suffix matching =="
python - <<'PY'
examples = [
("api.openai.com", "openai.com"),
("myopenai.com", "openai.com"),
("api.x.ai", "x.ai"),
("teamx.ai", "x.ai"),
("generativelanguage.googleapis.com", "googleapis.com"),
("notgoogleapis.com", "googleapis.com"),
]
for host, suffix in examples:
print(f"{host:35} endswith({suffix}) = {host.endswith(suffix)}")
PYRepository: lightseekorg/smg
Length of output: 1629
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Check usage of WorkerKind::External after classification =="
rg -n 'WorkerKind::External' --type rust -B2 -A5 | head -60
echo
echo "== Check if there are other validations in the classify function =="
cat -n model_gateway/src/core/steps/worker/classify.rs | head -150
echo
echo "== Look for tests that might cover provider detection =="
fd -0 'test|spec' | xargs -0 rg -l 'ProviderType|from_url' | head -10Repository: lightseekorg/smg
Length of output: 10119
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Look for tests in crates/protocols for ProviderType =="
find . -name '*.rs' -path '*/test*' -o -name '*_test.rs' -o -name 'tests.rs' | xargs grep -l 'ProviderType\|from_url' 2>/dev/null || echo "No test files found"
echo
echo "== Search for any test modules in worker.rs =="
tail -100 crates/protocols/src/worker.rs | grep -A50 '#\[cfg(test)\]' || echo "No test module found in worker.rs"
echo
echo "== Check classify.rs for any unit tests =="
tail -100 model_gateway/src/core/steps/worker/classify.rs | grep -A50 '#\[cfg(test)\]' || echo "No test module found in classify.rs"
echo
echo "== Verify the actual fix needed: show the from_url function in full =="
sed -n '214,226p' crates/protocols/src/worker.rsRepository: lightseekorg/smg
Length of output: 843
Fix hostname boundary matching in ProviderType::from_url to prevent lookalike domain misclassification.
The provider detection at line 105 relies on ProviderType::from_url, which uses raw host.ends_with() checks. This allows lookalike domains such as myopenai.com, teamx.ai, and notgoogleapis.com to match as legitimate providers. When matched, workers skip all probing steps (health checks, gRPC, model validation) and are immediately classified as WorkerKind::External without further verification.
Update the helper in crates/protocols/src/worker.rs:214-226 to require an exact host match or a proper . boundary (e.g., host == "openai.com" || host.ends_with(".openai.com")), and add regression test cases covering boundary cases:
api.openai.com→ match (valid subdomain)myopenai.com→ no match (lookalike)api.x.ai→ match (valid subdomain)teamx.ai→ no match (lookalike)generativelanguage.googleapis.com→ match (valid subdomain)notgoogleapis.com→ no match (lookalike)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/core/steps/worker/classify.rs` around lines 105 - 111,
ProviderType::from_url currently uses host.ends_with(...) which false-positives
lookalike domains; change its matching logic in crates/protocols/src/worker.rs
(the ProviderType::from_url helper) to require exact host equality or a
dot-prefixed suffix (e.g., host == "openai.com" ||
host.ends_with(".openai.com")) for each provider, and update all provider checks
(OpenAI, x.ai, googleapis/generative language) accordingly so subdomains match
but lookalikes like myopenai.com or notgoogleapis.com do not. Also add
regression tests covering the specified boundary cases (api.openai.com → match,
myopenai.com → no match, api.x.ai → match, teamx.ai → no match,
generativelanguage.googleapis.com → match, notgoogleapis.com → no match) to
validate the fix.
Add TestIGWMixedWorkerClassification that exercises the classify step with a mix of local sglang backends and external cloud providers, verifying race-condition-free classification. Test setup (requires 4 GPUs + OPENAI_API_KEY + XAI_API_KEY): - 2 HTTP sglang workers (Llama-3.1-8B-Instruct) - 2 gRPC sglang workers (DeepSeek-R1-Distill-Qwen-7B) - 1 OpenAI external worker (https://api.openai.com) - 1 xAI external worker (https://api.x.ai) Gateway starts in IGW mode first, then all workers are added immediately without waiting for registration — this exercises the startup race condition that previously caused misclassification. Assertions: - All 6 workers register successfully - Local workers have runtime_type sglang/vllm/trtllm (not external) - Cloud workers have runtime_type external - /v1/models returns models from all workers (local + external) Signed-off-by: Simon Lin <simonslin@gmail.com> Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f6c1362e2a
ℹ️ 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".
| if let Some(provider) = ProviderType::from_url(&config.url) { | ||
| debug!( | ||
| "Worker {} explicitly configured as {} → Local", | ||
| config.url, config.runtime_type | ||
| "Worker {} URL matches known provider ({}) → External", | ||
| config.url, provider | ||
| ); | ||
| context.data.worker_kind = Some(WorkerKind::Local); | ||
| context.data.worker_kind = Some(WorkerKind::External); | ||
| return Ok(StepResult::Success); |
There was a problem hiding this comment.
This change makes Unspecified workers external only when ProviderType::from_url matches a hard-coded host, so any other OpenAI-compatible endpoint (custom domain, self-hosted proxy, third-party provider) now falls through to Local classification and the workflow enters DetectConnectionModeStep, which only probes /health and gRPC for local workers and can fail registration entirely. Before this commit, a reachable /v1/models response was enough to classify these endpoints as external, so this is a behavioral regression for existing callers that rely on default runtime detection.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@e2e_test/router/test_worker_api.py`:
- Around line 438-443: The assertions use redundant case checks; simplify by
converting each model string to lower-case once and test for the lowercase
substrings. Update the assertions referencing model_ids so they use a single
case-insensitive check like "llama" in m.lower() and "deepseek" in m.lower()
(remove the redundant 'Llama'/'DeepSeek' checks) to make the intent clearer and
avoid duplicate conditions.
- Around line 411-413: Replace the hard-coded counts with derived counts
computed from the actual worker collections: compute expected_local_count from
the list/criteria that identifies local workers (e.g., count workers with HTTP
or gRPC protocol) and compute expected_total_count by adding OpenAI/xAI entries
or by taking len(workers_list_sources) instead of using 4/6; then update the
assert currently referencing workers and the message string to use these derived
values (keep the asserted object as workers and the message to reflect the
computed expected counts). Ensure you update any helper variables used in the
test (the variables that identify local workers or external workers) so the
assertion remains self-documenting and resilient to changes.
- Around line 386-405: The healthy_local count currently sums all healthy
workers from gateway.list_workers(), which includes external workers added with
disable_health_check=True; update the loop to count only local workers by
filtering workers by their URL or by comparing against a stored set/list of
local worker URLs used when launching them (use gateway.list_workers() and
filter by w.url in local_urls or by a prefix identifying local addresses), then
compute healthy_local = sum(1 for w in workers if w.status == "healthy" and
w.url in local_urls) so the loop waits for 4 healthy local workers rather than
any 4 healthy workers.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: aa2afabd-6e62-4e6d-8623-578da294b65e
📒 Files selected for processing (1)
e2e_test/router/test_worker_api.py
Add a 4-GPU CI job that runs gateway tests requiring multiple GPUs. This enables the TestIGWMixedWorkerClassification test which needs 4 GPUs (2 HTTP + 2 gRPC sglang workers) plus external API keys (OPENAI_API_KEY, XAI_API_KEY) to verify worker classification. What changed: - pr-test-rust.yml: add e2e-4gpu-gateway job using e2e-gpu-job.yml reusable workflow with sglang engine, 4-gpu-h100 runner, and test filter for gpu(4) tests in e2e_test/router - pr-test-rust.yml: add e2e-4gpu-gateway to finish job dependencies Note: the test will be skipped on runners without OPENAI_API_KEY and XAI_API_KEY environment variables (pytest.skip in the test itself). Signed-off-by: Simon Lin <simonslin@gmail.com> Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 89ee206c05
ℹ️ 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".
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
e2e_test/router/test_worker_api.py (2)
386-407:⚠️ Potential issue | 🟠 MajorBug:
healthy_localcounts all healthy workers, including external ones.External workers added with
disable_health_check=Truebecome immediately healthy. The loop could exit early when 2 external + 2 local workers are healthy (total 4), before all 4 local workers are ready.🐛 Proposed fix to filter by local worker URLs
+ local_urls = {w.base_url for w in all_local_workers} + # Wait for local workers to be registered (external are instant with disable_health_check) deadline = time.perf_counter() + 120 while time.perf_counter() < deadline: workers = gateway.list_workers() - healthy_local = sum(1 for w in workers if w.status == "healthy") + healthy_local = sum( + 1 for w in workers + if w.status == "healthy" and w.url in local_urls + ) if healthy_local >= 4: # 2 HTTP + 2 gRPC break time.sleep(2)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@e2e_test/router/test_worker_api.py` around lines 386 - 407, The loop currently sets healthy_local = sum(1 for w in workers if w.status == "healthy") which counts external workers too; change this to only count local workers by filtering worker URLs (e.g., only include w where w.status == "healthy" and w.url indicates a local address such as starting with "http://127.0.0.1", "http://localhost", "grpc://127.0.0.1", etc.) when computing healthy_local and when formatting the pytest.fail message so the timeout check requires 4 healthy local workers (use gateway.list_workers() and the w.url filter in both the loop condition and the failure diagnostic).
438-443: 🧹 Nitpick | 🔵 TrivialRedundant case-insensitive checks.
"llama" in m.lower()already handles all case variations including "Llama". Theor "Llama" in mpart is redundant. Same applies to the DeepSeek check.♻️ Simplified assertions
- assert any("llama" in m.lower() or "Llama" in m for m in model_ids), ( + assert any("llama" in m.lower() for m in model_ids), ( f"Expected a Llama model from HTTP workers, got: {model_ids}" ) - assert any("deepseek" in m.lower() or "DeepSeek" in m for m in model_ids), ( + assert any("deepseek" in m.lower() for m in model_ids), ( f"Expected a DeepSeek model from gRPC workers, got: {model_ids}" )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@e2e_test/router/test_worker_api.py` around lines 438 - 443, The assertions in test_worker_api.py use redundant case checks; simplify them to use a single case-insensitive test on model_ids by calling m.lower() once. Replace the expressions checking for Llama and DeepSeek with any("llama" in m.lower() for m in model_ids) and any("deepseek" in m.lower() for m in model_ids) respectively so the variable model_ids is checked case-insensitively without the unnecessary OR branches.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In @.github/workflows/pr-test-rust.yml:
- Around line 582-593: The pytest -k filter in the e2e-4gpu-gateway job is
ineffective because "gpu(4)" is not present in test node IDs; update the job's
test_filter (in the e2e-4gpu-gateway job) to either just "-k
'TestIGWMixedWorkerClassification'" or remove test_filter entirely and rely on
the existing E2E_GPU_TIER="4" + pytest_collection_modifyitems hook to select
4-GPU tests, ensuring the filter matches actual test node names rather than
marker arguments.
---
Duplicate comments:
In `@e2e_test/router/test_worker_api.py`:
- Around line 386-407: The loop currently sets healthy_local = sum(1 for w in
workers if w.status == "healthy") which counts external workers too; change this
to only count local workers by filtering worker URLs (e.g., only include w where
w.status == "healthy" and w.url indicates a local address such as starting with
"http://127.0.0.1", "http://localhost", "grpc://127.0.0.1", etc.) when computing
healthy_local and when formatting the pytest.fail message so the timeout check
requires 4 healthy local workers (use gateway.list_workers() and the w.url
filter in both the loop condition and the failure diagnostic).
- Around line 438-443: The assertions in test_worker_api.py use redundant case
checks; simplify them to use a single case-insensitive test on model_ids by
calling m.lower() once. Replace the expressions checking for Llama and DeepSeek
with any("llama" in m.lower() for m in model_ids) and any("deepseek" in
m.lower() for m in model_ids) respectively so the variable model_ids is checked
case-insensitively without the unnecessary OR branches.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: ae1318b0-ccd0-4774-9530-e6a0c1f4cd74
📒 Files selected for processing (2)
.github/workflows/pr-test-rust.ymle2e_test/router/test_worker_api.py
The -k flag uses Python expressions, not marker syntax. gpu(4) is not valid — use the class name directly. Signed-off-by: Simon Lin <simonslin@gmail.com> Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 42f27504b2
ℹ️ 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".
| json={ | ||
| "url": "https://api.openai.com", | ||
| "api_key": openai_key, | ||
| "runtime": "external", |
There was a problem hiding this comment.
TestIGWMixedWorkerClassification claims to validate URL-based external classification, but this request sets "runtime": "external", which makes ClassifyWorkerTypeStep short-circuit on explicit runtime and skip URL detection entirely. As a result, this new gated test can pass even if ProviderType::from_url classification is broken, so it does not protect the regression this commit is intended to cover.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
♻️ Duplicate comments (1)
model_gateway/src/core/steps/worker/classify.rs (1)
104-112:⚠️ Potential issue | 🟠 MajorHostname boundary matching issue persists.
The past review comment regarding
ProviderType::from_urlusing rawends_with()checks (allowing lookalike domains likemyopenai.comto match) remains unaddressed. This allows malicious workers to bypass health/gRPC/models probing by using lookalike domains and being immediately classified as External.The fix must be applied in
crates/protocols/src/worker.rs(theProviderType::from_urlimplementation) to require exact host match or proper subdomain boundary (e.g.,host == "openai.com" || host.ends_with(".openai.com")), with regression tests for boundary cases.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/steps/worker/classify.rs` around lines 104 - 112, ProviderType::from_url currently uses raw ends_with checks causing lookalike hosts (e.g., myopenai.com) to be treated as provider domains and classify workers as WorkerKind::External; update the implementation in crates/protocols/src/worker.rs (ProviderType::from_url) to parse the URL host and require either host == "openai.com" or host.ends_with(".openai.com") (and do the same for other provider domains) so only exact hosts or valid subdomains match, and add regression tests covering boundary cases like "myopenai.com", "openai.com", "api.openai.com", and "bad.openai.com.evil" to ensure correct behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@model_gateway/src/core/steps/worker/classify.rs`:
- Around line 104-112: ProviderType::from_url currently uses raw ends_with
checks causing lookalike hosts (e.g., myopenai.com) to be treated as provider
domains and classify workers as WorkerKind::External; update the implementation
in crates/protocols/src/worker.rs (ProviderType::from_url) to parse the URL host
and require either host == "openai.com" or host.ends_with(".openai.com") (and do
the same for other provider domains) so only exact hosts or valid subdomains
match, and add regression tests covering boundary cases like "myopenai.com",
"openai.com", "api.openai.com", and "bad.openai.com.evil" to ensure correct
behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: b502953b-b5dd-4c18-bd94-72f4382f7133
📒 Files selected for processing (1)
model_gateway/src/core/steps/worker/classify.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6ec46bac45
ℹ️ 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".
6ec46ba to
e3f1c86
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e3f1c8642b
ℹ️ 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".
e3f1c86 to
269f62d
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 269f62d410
ℹ️ 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".
| // Explicit local runtime (Sglang, Vllm, Trtllm) → Local | ||
| if config.runtime_type.is_specified() { | ||
| // 3. URL matches known cloud provider → External (no probing needed) | ||
| if let Some(provider) = ProviderType::from_url(&config.url) { |
There was a problem hiding this comment.
Fresh evidence in the current tree shows ProviderType::from_url still matches hosts with ends_with(...) (crates/protocols/src/worker.rs:217-224), so this new early-return path can classify domains like notopenai.com as OpenAI and force External before /health/gRPC probes run. In practice, an unspecified local worker on a similarly suffixed hostname is routed down the external branch and can fail registration or be handled by the wrong workflow.
Useful? React with 👍 / 👎.
| json={ | ||
| "url": "https://api.openai.com", | ||
| "api_key": openai_key, | ||
| "runtime": "external", |
There was a problem hiding this comment.
Fresh evidence: the classifier still short-circuits on explicit runtime (classify.rs:90-102), and this test payload sets "runtime": "external", so the new 4-GPU gate does not actually exercise URL-based auto-classification for cloud workers. That means regressions in the URL detection path can pass this job even though the test is intended to validate that behavior.
Useful? React with 👍 / 👎.
269f62d to
a49a827
Compare
Add TestIGWMixedWorkerClassification that exercises the classify step with a mix of local sglang backends and external cloud providers, verifying race-condition-free classification. Test setup (requires 4 GPUs + OPENAI_API_KEY + XAI_API_KEY): - 2 HTTP sglang workers (Llama-3.1-8B-Instruct) - 2 gRPC sglang workers (DeepSeek-R1-Distill-Qwen-7B) - 1 OpenAI external worker (https://api.openai.com) - 1 xAI external worker (https://api.x.ai) Gateway starts in IGW mode first, then all workers are added immediately without waiting for backends to be ready — this exercises the startup race condition that previously caused misclassification. Add wait_ready parameter to Worker.start() and start_workers() to support spawning backend processes without blocking on health checks. Defaults to True for backward compatibility. Assertions: - All 6 workers register successfully - Local workers have runtime_type sglang/vllm/trtllm (not external) - Cloud workers have runtime_type external - /v1/models returns models from all workers (local + external) Signed-off-by: Simon Lin <simonslin@gmail.com> Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
a49a827 to
f74502b
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f74502bdf7
ℹ️ 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".
| // Explicit local runtime (Sglang, Vllm, Trtllm) → Local | ||
| if config.runtime_type.is_specified() { | ||
| // 3. URL matches known cloud provider → External (no probing needed) | ||
| if let Some(provider) = ProviderType::from_url(&config.url) { |
There was a problem hiding this comment.
ProviderType::from_url(&config.url) now decides the known-provider fast path, but it only matches parseable absolute URLs. If a worker is added with an unspecified runtime and a scheme-less host (for example api.openai.com, which other worker URL helpers in this flow support by adding http://), this branch is skipped and the new logic falls through to the default-Local path instead of External. That can send a real cloud endpoint through local /health/gRPC registration and fail worker onboarding. Please normalize or reparse scheme-less URLs before this provider check.
Useful? React with 👍 / 👎.
1 participant
Summary
Fix the race condition in
ClassifyWorkerTypeStepwhere local backends starting up were misclassified as External.Refs: #820
What changed
Rewrote the auto-detection logic in
classify.rsto eliminate the unreliable/v1/modelsreachability heuristic:runtime_type/healthresponds/v1/modelsowned_bymatches sglang/vllm/trtllmWhy
The old heuristic was:
/v1/modelsresponds + no/health= External. This caused a race condition — local backends (sglang, vllm) can serve/v1/modelsbefore/healthis ready during startup, causing them to be misclassified as External. This is the root cause of the issue described in #820.The new approach uses two safer signals:
ProviderType::from_url()— instant, deterministic, zero race conditionsowned_byfield from/v1/modelsresponse — confirms local backends rather than inferring ExternalHow
is_models_endpoint_reachable()(status-code-only check)probe_models_owned_by()which parses the/v1/modelsresponse body and checks if the first model'sowned_byfield matches a known local backendProviderType::from_url()check before any probing — known cloud URLs are classified instantlyis_specified()checkTrade-off: External providers on private IPs (e.g.,
http://10.0.0.1:8080proxying to OpenAI) must setruntime_type: externalexplicitly. URL-based detection cannot identify them.Test plan
cargo check -p smg— compiles cleancargo test -p smg— all tests pass (0 failures)cargo clippy -p smg --all-targets --all-features -- -D warnings— no warningshttps://api.openai.comwith noruntime_type→ should classify as External via URL detectionruntime_typeduring startup → should classify as Local (not External)Summary by CodeRabbit
Bug Fixes
Tests
Chores