Repository navigation
Conversation
github-actions
Bot
added
tests
📝 WalkthroughWalkthroughThis PR implements model alias grouping and canonical model resolution across the system. It groups model variants (OpenAI date-suffixed and XAI Grok revisions) into ModelCard aliases, resolves those aliases to canonical upstream IDs at the router level, and updates worker/registry/selection logic to index and match against aliases. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Router as OpenAI Router
participant Worker as Worker<br/>(Model View)
participant Upstream as Upstream<br/>Provider
Client->>Router: Request<br/>(model: "grok-4")
Router->>Router: request_provider("grok-4")<br/>→ ProviderType
Router->>Worker: select_worker(model, provider)
Worker-->>Router: return Worker
Router->>Worker: resolve_effective_model<br/>(worker, "grok-4")
Worker->>Worker: find_model("grok-4")<br/>in aliases
Worker-->>Router: canonical_id: "grok-4-0709"
Router->>Router: resolve_provider<br/>(worker, "grok-4-0709")
Router->>Router: Patch request model<br/>to "grok-4-0709"
Router->>Upstream: Forward request<br/>(model: "grok-4-0709")
Upstream-->>Client: Response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested labels
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)
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 |
|
Hi @khoaatra, the DCO sign-off check has failed. All commits must include a To fix existing commits: # Sign off the last N commits (replace N with the number of unsigned commits)
git rebase HEAD~N --signoff
git push --force-with-leaseTo sign off future commits automatically:
|
![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 introduces model alias support and canonical ID resolution, primarily for OpenAI and xAI providers. It implements logic to group discovered models into cards, indexes these aliases in the worker registry, and ensures that requests are rewritten to use canonical IDs before being forwarded. Review feedback identifies critical performance bottlenecks on hot routing paths caused by frequent cloning of model lists and inefficient merging logic, recommending the implementation of more direct query methods for model support and provider information.
| self.effective_models() | ||
| .into_iter() | ||
| .any(|model| model.matches(model_id)) |
There was a problem hiding this comment.
Rebuilding the entire merged model list via effective_models() on every call to supports_model is a significant performance regression. This function is on the hot path for request routing and now involves multiple allocations, clones, and an O(N*M) merge loop. It is much more efficient to check the overridden models and metadata models separately using their existing supports methods. This approach prioritizes code simplicity and avoids unnecessary overhead on hot paths.
if overridden.supports(model_id) {
return true;
}
self.metadata.supports_model(model_id)References
- Prioritize code simplicity and clarity over micro-optimizations, especially when the performance gain is negligible for typical use cases.
| worker | ||
| .models() | ||
| .into_iter() | ||
| .find(|candidate| candidate.matches(model)) | ||
| .map(|candidate| candidate.id) |
There was a problem hiding this comment.
Calling worker.models().into_iter().find(...) on every request is a major performance bottleneck. worker.models() clones the entire list of model cards, which can be large for external providers. This overhead will significantly impact request latency. The Worker trait should be extended with methods to resolve canonical IDs and providers efficiently to avoid full list clones on the hot path. This optimization should be extracted into a shared mechanism to ensure consistency across all routing paths.
References
- If an optimization is applicable to multiple code paths, extract it into a shared helper function to ensure consistency and avoid code duplication.
| if let Some(pt) = worker | ||
| .models() | ||
| .into_iter() | ||
| .find(|candidate| candidate.matches(model)) | ||
| .and_then(|candidate| candidate.provider) | ||
| { |
There was a problem hiding this comment.
Similar to resolve_effective_model, this block calls worker.models(), causing expensive full-list clones on every request. This logic should be moved into a more efficient method on the Worker trait or BasicWorker implementation that can query the provider for a model without cloning the entire model list. Applying this optimization consistently across code paths aligns with repository guidelines.
References
- If an optimization is applicable to multiple code paths, extract it into a shared helper function to ensure consistency and avoid code duplication.
| fn worker_provider_hint(worker: &Arc<dyn Worker>) -> Option<ProviderType> { | ||
| worker | ||
| .default_provider() | ||
| .cloned() | ||
| .or_else(|| worker.models().into_iter().find_map(|model| model.provider)) | ||
| .or_else(|| ProviderType::from_url(worker.url())) | ||
| } | ||
|
|
||
| fn worker_provider_for_request(worker: &Arc<dyn Worker>, model_id: &str) -> Option<ProviderType> { | ||
| worker | ||
| .models() | ||
| .into_iter() | ||
| .find(|model| model.matches(model_id)) | ||
| .and_then(|model| model.provider) | ||
| .or_else(|| worker.provider_for_model(model_id).cloned()) | ||
| .or_else(|| ProviderType::from_model_name(model_id)) | ||
| .or_else(|| worker_provider_hint(worker)) | ||
| } |
There was a problem hiding this comment.
Both worker_provider_hint and worker_provider_for_request call worker.models(), which clones the entire model list for every candidate worker during selection. In a system with many workers and models, this will cause massive CPU and memory pressure during routing. These should be optimized by providing a direct way to query provider information from the worker without cloning the entire model list, ensuring this optimization is shared across all applicable paths.
References
- If an optimization is applicable to multiple code paths, extract it into a shared helper function to ensure consistency and avoid code duplication.
![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: 957e950678
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let providers: HashSet<_> = workers | ||
| .iter() | ||
| .filter_map(|worker| worker_provider_hint(worker)) | ||
| .collect(); | ||
| let has_multiple_providers = providers.len() > 1; |
There was a problem hiding this comment.
filter_by_provider currently builds the provider set with filter_map, which drops workers that have no provider hint. In a mixed deployment (for example, one Anthropic worker plus one provider-unspecified worker), providers.len() becomes 1, so filtering is disabled and the unspecified worker remains eligible for Anthropic/XAI requests. That breaks the intended credential-isolation behavior in this module and can route provider-specific API keys to the wrong backend.
Useful? React with 👍 / 👎.
| .find(|model| model.matches(model_id)) | ||
| .and_then(|model| model.provider) | ||
| .or_else(|| worker.provider_for_model(model_id).cloned()) | ||
| .or_else(|| ProviderType::from_model_name(model_id)) |
There was a problem hiding this comment.
worker_provider_for_request falls back to ProviderType::from_model_name(model_id), which uses the requested model string rather than worker metadata. A wildcard or provider-unspecified worker can therefore be classified as matching any provider implied by the request (e.g., claude-*), bypassing provider isolation when filtering is active and allowing misrouting/credential leakage to non-target workers.
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: 5
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/routers/openai/chat.rs (1)
61-66:⚠️ Potential issue | 🟠 MajorDo not hard-filter provider before alias resolution.
request_provider(model)is evaluated on the raw request string here, but alias canonicalization only happens after a worker is selected. Becauserequest_providerfalls back to OpenAI for unknown names, any alias that exists only in discoveredModelCard.aliaseswill be filtered out before refresh/selection in a mixed-provider setup, so the request never reaches the correct worker.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/openai/chat.rs` around lines 61 - 66, The code is prematurely setting the provider field using request_provider(model) when calling selector.select_worker(&SelectWorkerRequest { ... }), which filters out workers before alias canonicalization (ModelCard.aliases) runs; remove or set provider to None when constructing SelectWorkerRequest in chat.rs so worker selection can consider aliases, and instead apply request_provider(model) only after a worker is selected (or after alias canonicalization) to determine the final provider. Ensure changes target the selector.select_worker call and the SelectWorkerRequest construction to avoid hard-filtering by provider.
🤖 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/routers/openai/responses/route.rs`:
- Around line 63-67: The provider hint is being applied too early on the
SelectWorkerRequest causing premature filtering; remove or omit the provider:
Some(request_provider(model)) from the .select_worker(...) call (the
SelectWorkerRequest construction) so alias-based models in discovered metadata
aren't excluded before alias resolution, and instead pass the provider hint only
after resolve_effective_model(...) has run and returned the resolved
model/provider.
In `@model_gateway/src/routers/openai/router.rs`:
- Around line 43-50: The realtime REST/WS handlers still pass the raw model name
into upstream proxies, so alias routing breaks; update each call site that
forwards realtime requests (calls to forward_realtime_rest and
handle_realtime_ws) to first canonicalize the model using
resolve_effective_model(worker, model) and pass that returned String into the
proxy functions instead of the original model variable; ensure both REST and WS
branches (the locations that currently call forward_realtime_rest(...) and
handle_realtime_ws(...)) are changed so the canonical ID is used for upstream
forwarding.
In `@model_gateway/src/routers/worker_selection.rs`:
- Around line 160-163: The provider filter currently ignores None when deciding
if multiple provider-buckets exist, causing mixed pools like {Some(XAI), None}
to be treated as single-provider and improperly filtered; update the logic that
counts distinct providers (the code path using has_multiple_providers or the
provider-counting routine) to treat None as its own distinct provider bucket
(i.e., include Option::None as a key when building the provider set) so that
external_workers.retain(|w| worker_matches_provider(w, p, None)) and the refresh
branches consistently handle None-workers; apply the same change to the other
occurrences noted (the blocks around lines 189-205 and 229-237) so
provider-aware filtering and refreshing treat None as a distinct provider rather
than ignoring it.
In `@model_gateway/src/worker/worker.rs`:
- Around line 401-407: canonical_model_id currently resolves only against
metadata.spec.models (via metadata().canonical_model_id(model_id)) which is out
of sync after set_models() updates aliases; update canonical_model_id to consult
the dynamic effective model set instead. Specifically, change canonical_model_id
to ask effective_models() (or delegate to a new metadata wrapper that computes
canonical ids from effective_models()) so it returns the discovered upstream ID
for aliases added by set_models(); mirror the same fix for the other occurrences
noted (around the supports_model()/models() call sites and the blocks at the
other lines referenced) so canonicalization and routing use the same
effective_models() source.
In `@model_gateway/src/workflow/steps/external/discover_models.rs`:
- Around line 64-121: The alias-canonicalization and primary-selection logic in
alias_group_key, xai_revision_rank, and select_primary_and_aliases is duplicated
and must be centralized; remove these local implementations and call the
canonical helpers from the protocols module (the same functions used in
crates::protocols::models), or extract them into a shared helper crate and
import that helper here; update imports to use the shared functions (e.g.,
alias_group_key, xai_revision_rank, select_primary_and_aliases or appropriately
named equivalents) and replace local calls and references so discovery uses the
single canonical implementation used by routing.
---
Outside diff comments:
In `@model_gateway/src/routers/openai/chat.rs`:
- Around line 61-66: The code is prematurely setting the provider field using
request_provider(model) when calling selector.select_worker(&SelectWorkerRequest
{ ... }), which filters out workers before alias canonicalization
(ModelCard.aliases) runs; remove or set provider to None when constructing
SelectWorkerRequest in chat.rs so worker selection can consider aliases, and
instead apply request_provider(model) only after a worker is selected (or after
alias canonicalization) to determine the final provider. Ensure changes target
the selector.select_worker call and the SelectWorkerRequest construction to
avoid hard-filtering by provider.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: e73bc81d-e2c6-401d-bcd5-fe6c68b0179d
📒 Files selected for processing (9)
crates/protocols/src/models.rsmodel_gateway/src/routers/openai/chat.rsmodel_gateway/src/routers/openai/responses/route.rsmodel_gateway/src/routers/openai/router.rsmodel_gateway/src/routers/worker_selection.rsmodel_gateway/src/worker/registry.rsmodel_gateway/src/worker/worker.rsmodel_gateway/src/workflow/steps/external/discover_models.rsmodel_gateway/tests/routing/test_openai_routing.rs
| .select_worker(&SelectWorkerRequest { | ||
| model_id: model, | ||
| headers, | ||
| provider: Some(ProviderType::OpenAI), | ||
| provider: Some(request_provider(model)), | ||
| ..Default::default() |
There was a problem hiding this comment.
This provider hint is being applied too early.
Like the chat path, this filters on request_provider(model) before alias resolution. If the client sends an alias that only exists in discovered model metadata, the selector can exclude the real worker before resolve_effective_model(...) ever runs.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/routers/openai/responses/route.rs` around lines 63 - 67,
The provider hint is being applied too early on the SelectWorkerRequest causing
premature filtering; remove or omit the provider: Some(request_provider(model))
from the .select_worker(...) call (the SelectWorkerRequest construction) so
alias-based models in discovered metadata aren't excluded before alias
resolution, and instead pass the provider hint only after
resolve_effective_model(...) has run and returned the resolved model/provider.
| /// Resolve the canonical upstream model ID from the worker's live model view. | ||
| pub(super) fn resolve_effective_model(worker: &dyn Worker, model: &str) -> String { | ||
| worker | ||
| .models() | ||
| .into_iter() | ||
| .find(|candidate| candidate.matches(model)) | ||
| .map(|candidate| candidate.id) | ||
| .unwrap_or_else(|| worker.canonical_model_id(model).to_string()) |
There was a problem hiding this comment.
Canonicalize realtime model IDs before proxying.
This helper only fixes the call sites that use it. In this file, the realtime REST/WS paths still pass the raw model into forward_realtime_rest / handle_realtime_ws (Lines 206-214, 226-234, 245-253, and 272-278), so grok-4 can now select the xAI worker but still gets forwarded upstream as grok-4 instead of the worker’s canonical ID. That leaves alias routing broken for the realtime endpoints.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/routers/openai/router.rs` around lines 43 - 50, The
realtime REST/WS handlers still pass the raw model name into upstream proxies,
so alias routing breaks; update each call site that forwards realtime requests
(calls to forward_realtime_rest and handle_realtime_ws) to first canonicalize
the model using resolve_effective_model(worker, model) and pass that returned
String into the proxy functions instead of the original model variable; ensure
both REST and WS branches (the locations that currently call
forward_realtime_rest(...) and handle_realtime_ws(...)) are changed so the
canonical ID is used for upstream forwarding.
| fn alias_group_key(id: &str) -> String { | ||
| if let Some(captures) = XAI_REVISION_SUFFIX_PATTERN.captures(id) { | ||
| if let Some(base) = captures.get(1) { | ||
| return base.as_str().to_string(); | ||
| } | ||
| } | ||
|
|
||
| DATE_SUFFIX_PATTERN.replace(id, "").to_string() | ||
| } | ||
|
|
||
| fn xai_revision_rank(id: &str) -> Option<u16> { | ||
| XAI_REVISION_SUFFIX_PATTERN | ||
| .captures(id) | ||
| .and_then(|captures| captures.get(2)) | ||
| .and_then(|suffix| suffix.as_str().parse::<u16>().ok()) | ||
| } | ||
|
|
||
| fn select_primary_and_aliases(group_key: &str, variants: Vec<ModelInfo>) -> (String, Vec<String>) { | ||
| let primary_id = if variants.iter().any(|variant| variant.id == group_key) { | ||
| group_key.to_string() | ||
| } else if variants | ||
| .iter() | ||
| .all(|variant| xai_revision_rank(&variant.id).is_some()) | ||
| { | ||
| variants | ||
| .iter() | ||
| .max_by_key(|variant| { | ||
| ( | ||
| variant.created.unwrap_or(0), | ||
| xai_revision_rank(&variant.id).unwrap_or(0), | ||
| variant.id.as_str(), | ||
| ) | ||
| }) | ||
| .map(|variant| variant.id.clone()) | ||
| .unwrap_or_else(|| group_key.to_string()) | ||
| } else { | ||
| variants | ||
| .iter() | ||
| .map(|variant| variant.id.as_str()) | ||
| .min_by(|a, b| a.len().cmp(&b.len()).then_with(|| a.cmp(b))) | ||
| .map(ToOwned::to_owned) | ||
| .unwrap_or_else(|| group_key.to_string()) | ||
| }; | ||
|
|
||
| let mut aliases: Vec<String> = variants | ||
| .iter() | ||
| .filter_map(|variant| (variant.id != primary_id).then(|| variant.id.clone())) | ||
| .collect(); | ||
|
|
||
| if group_key != primary_id && !aliases.iter().any(|alias| alias == group_key) { | ||
| aliases.push(group_key.to_string()); | ||
| } | ||
|
|
||
| aliases.sort(); | ||
| aliases.dedup(); | ||
|
|
||
| (primary_id, aliases) | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Deduplicate the alias-canonicalization rules.
This block is now a second copy of the grouping/primary-selection logic already added in crates/protocols/src/models.rs:87-203. If one side changes without the other, discovery and routing will derive different canonical IDs/aliases and the mismatch will be silent. Please centralize this in a shared helper or delegate to the protocols implementation from here.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/workflow/steps/external/discover_models.rs` around lines 64
- 121, The alias-canonicalization and primary-selection logic in
alias_group_key, xai_revision_rank, and select_primary_and_aliases is duplicated
and must be centralized; remove these local implementations and call the
canonical helpers from the protocols module (the same functions used in
crates::protocols::models), or extract them into a shared helper crate and
import that helper here; update imports to use the shared functions (e.g.,
alias_group_key, xai_revision_rank, select_primary_and_aliases or appropriately
named equivalents) and replace local calls and references so discovery uses the
single canonical implementation used by routing.
224d701 to
c939d84
Compare
Signed-off-by: Khoa Tran <khoa.k.tran@oracle.com>
c939d84 to
cd3c16c
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cd3c16c4ba
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| } else { | ||
| merged_models.push(metadata_model.clone()); |
There was a problem hiding this comment.
When both models_override and metadata models are non-wildcard, this branch re-adds metadata-only entries that were not returned by discovery. In deployments that rely on lazy discovery to reflect current upstream availability (e.g., revoked or removed models), supports_model will still return true for stale metadata IDs and the router can continue forwarding requests for models the upstream no longer serves.
Useful? React with 👍 / 👎.
| self.effective_models() | ||
| .into_iter() | ||
| .any(|model| model.matches(model_id)) |
There was a problem hiding this comment.
supports_model now calls effective_models(), which materializes and clones the full model-card list before every membership check. This method is on the hot path during worker selection (called across candidate workers per request), so in environments with many discovered models/aliases it introduces avoidable allocation and copy overhead that can noticeably increase routing latency.
Useful? React with 👍 / 👎.
khoaatra
changed the title
1 participant
Description
Problem
Solution
Changes
Test Plan
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
New Features
Tests