fix(router): support model aliases end to end - #1964
Conversation
📝 WalkthroughWalkthroughModel aliases are indexed and resolved by ChangesModel Alias Routing
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant Router
participant WorkerRegistry
participant Worker
Client->>Router: request using model alias
Router->>WorkerRegistry: resolve_model_alias(alias)
WorkerRegistry-->>Router: canonical model ID
Router->>Worker: forward request with canonical model
Worker-->>Router: response
Router-->>Client: response
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
github-actions
Bot
added
documentation
9b65970 to
298a31f
Compare
github-actions
Bot
removed
the
documentation
junliu-mde
changed the title
7a3ee6c to
210de10
Compare
junliu-mde
changed the title
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![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: 210de10c16
ℹ️ 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".
| let canonical_model = self.worker_registry.resolve_model_alias(model_id); | ||
| let model_id = canonical_model.as_deref().unwrap_or(model_id); |
There was a problem hiding this comment.
When /v1/audio/transcriptions uses an alias, this canonicalization never runs because route_multipart_transcription bypasses route_typed_request. The alias-aware registry now finds the correct worker, but build_transcription_form still emits body.model at line 1088, so the backend receives the alias rather than the canonical model it serves and may reject the request. Resolve the alias on the multipart entry path and use that ID in the form.
Useful? React with 👍 / 👎.
| self.model_alias_index | ||
| .get(model_id) | ||
| .map(|workers| Arc::clone(&workers)) | ||
| .and_then(|canonical_id| { | ||
| self.model_index | ||
| .get(canonical_id.as_ref()) | ||
| .map(|workers| Arc::clone(&workers)) |
There was a problem hiding this comment.
With IGW enabled and an external worker declaring an alias, this lookup lets RouterManager select OpenAIRouter, but that router never resolves the alias: for example, routers/openai/chat.rs patches the outbound payload with the original model_id at line 108, and the Responses and realtime paths likewise forward the client name. An upstream provider that only serves the canonical model therefore rejects an alias that the gateway now claims to route; canonicalize at the OpenAI router entry as done for the other routers.
Useful? React with 👍 / 👎.
| if let Some(canonical_model) = canonical_model { | ||
| super::set_request_model(&mut json_val, canonical_model); | ||
| } |
There was a problem hiding this comment.
For an aliased /v1/rerank request, this rewrite sends the canonical ID upstream but leaves the original typed request unchanged. After the worker responds, build_rerank_response constructs the public response using req.model.clone() at line 1056, so this endpoint reports the alias while the other newly canonicalized responses report the model actually run. Thread the canonical ID into the rerank response construction as well.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
210de10 to
ee39fcf
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ee39fcfc17
ℹ️ 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".
| self.model_index | ||
| .get(canonical_id.as_ref()) | ||
| .map(|workers| Arc::clone(&workers)) |
There was a problem hiding this comment.
When IGW has a self-hosted worker advertising alias A for canonical model C and an external-provider worker serving C but not A (a common hybrid-migration setup), this expands A to every worker indexed under C. RouterManager::select_router_for_workers then prioritizes the unrelated external worker and selects its provider router; that router filters candidates with supports_model(A), so the external worker is rejected and provider filtering can exclude the self-hosted worker, returning model-not-found instead of serving the advertised alias. Preserve the canonical expansion for backend dispatch, but restrict router-selection candidates to workers that support the client-supplied alias.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@model_gateway/src/routers/http/router.rs`:
- Around line 250-255: Add a shared WorkerRegistry helper for canonicalizing a
model ID, returning the resolved canonical value or the original input when no
alias exists, with a lifetime-safe return type. Replace the inline
resolve-and-shadow logic at model_gateway/src/routers/http/router.rs lines
250-255 and 547-553, and model_gateway/src/routers/http/pd_router.rs lines
293-299, with this helper while preserving each call site’s existing effective
model ID behavior.
🪄 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 Plus
Run ID: c1561a2b-ffcb-4d5f-8b88-d3b79c280074
📒 Files selected for processing (14)
model_gateway/src/routers/grpc/common/responses/utils.rsmodel_gateway/src/routers/grpc/common/stages/dispatch_metadata.rsmodel_gateway/src/routers/grpc/context.rsmodel_gateway/src/routers/grpc/pipeline.rsmodel_gateway/src/routers/grpc/router.rsmodel_gateway/src/routers/http/mod.rsmodel_gateway/src/routers/http/pd_router.rsmodel_gateway/src/routers/http/router.rsmodel_gateway/src/routers/router_manager.rsmodel_gateway/src/worker/registry.rsmodel_gateway/tests/common/mock_worker.rsmodel_gateway/tests/routing/mod.rsmodel_gateway/tests/routing/model_alias_test.rsmodel_gateway/tests/routing/pd_routing_test.rs
A worker registers under a canonical model ID and may declare aliases next to it. Routing only ever matched the canonical ID, so a request that named an alias found no workers and got a 404. The registry now keeps an alias index next to the model index. It is separate so aliases stay out of `/v1/models` and out of the statistics, and `get_by_model` falls back to it. Every other per-model map — hash rings, retry overrides, load balancing policies — stays keyed by the canonical ID, so each request entry point resolves the alias exactly once and passes the canonical ID from there on: - the gRPC pipeline resolves in `RequestContext::new` and rewrites both `input.model_id` and the request's own `model` field, so worker selection, tokenizer lookup, parser selection and tool call ID format all read the canonical ID without resolving anything themselves; - both HTTP routers resolve at the top of their request path and rewrite the `model` field of the body they forward. The backend was registered under the canonical ID and has never heard of the alias. One visible consequence: the response reports the canonical model ID, not the alias the client sent, matching how the OpenAI API answers with the model it actually ran. `/v1/responses` gates on the client-supplied name before the pipeline canonicalizes it, so it gets `WorkerRegistry::contains_model`, which accepts both spellings. The `unknown` wildcard is not a registered name and stays rejected. Alias conflicts are resolved so that a name is either a canonical model ID or an alias, never both: - an alias naming a registered model is refused, and an alias already recorded is dropped when a model later claims that name — otherwise the name would start resolving to a different model once the real one's last worker left; - when two models declare the same alias the first registration wins, and losing that worker hands the alias to a remaining model that declares it rather than stranding it; - the alias entry is held across that whole decision, so a concurrent registration cannot be erased between the check and the delete. Per-model retry overrides are keyed by the canonical ID and are read before the pipeline canonicalizes, so `GrpcRouter::resolve_retry_config` resolves the alias itself. Without it a request naming an alias silently took the router default instead of the worker's override, on chat, generate, messages and completion alike. Tests cover the registry rules, both HTTP request paths end to end against the body the worker actually received, the multipart form, the rerank response, the gRPC pipeline, the Responses canonicalization, the retry override under an alias, and the Responses gate. Signed-off-by: Jun Liu <jun.c.liu@rakuten.com>
ee39fcf to
704ec6f
Compare
There was a problem hiding this comment.
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/http/router.rs (1)
1311-1336: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
route_rerankresolves the alias twice — a TOCTOU gap between the reported model and the one actually dispatched.
canonical_modelis resolved here at line 1318, beforeroute_typed_requestawaits the worker round-trip;route_typed_requestindependently re-resolves the same alias internally (line 254) at a later point. If the registry's alias mapping changes in between (worker replace/remove), the model namebuild_rerank_responsereports could diverge from the model that actually ran — the exact guaranteetest_rerank_forwards_and_reports_the_canonical_modelchecks for. Resolving right before the response is built (after the await, on the success path) shrinks this window to effectively nothing.🔧 Proposed fix: resolve after dispatch instead of before
async fn route_rerank( &self, headers: Option<&HeaderMap>, _tenant_meta: &TenantRequestMeta, body: &RerankRequest, model_id: &str, ) -> Response { - let canonical_model = self.worker_registry.resolve_model_alias(model_id); let response = self .route_typed_request(headers, body, "/v1/rerank", model_id) .await; if response.status().is_success() { + let canonical_model = self.worker_registry.resolve_model_alias(model_id); match Self::build_rerank_response(body, canonical_model.as_deref(), response).await {🤖 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 `@model_gateway/src/routers/http/router.rs` around lines 1311 - 1336, Move the worker_registry.resolve_model_alias(model_id) call in route_rerank to the successful response path after route_typed_request has awaited dispatch and immediately before build_rerank_response. Keep the resolved canonical model passed to build_rerank_response, ensuring the reported model reflects the latest alias mapping.
🤖 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.
Outside diff comments:
In `@model_gateway/src/routers/http/router.rs`:
- Around line 1311-1336: Move the worker_registry.resolve_model_alias(model_id)
call in route_rerank to the successful response path after route_typed_request
has awaited dispatch and immediately before build_rerank_response. Keep the
resolved canonical model passed to build_rerank_response, ensuring the reported
model reflects the latest alias mapping.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 3f2b98b6-994c-46be-ad55-4060735ad963
📒 Files selected for processing (14)
model_gateway/src/routers/grpc/common/responses/utils.rsmodel_gateway/src/routers/grpc/common/stages/dispatch_metadata.rsmodel_gateway/src/routers/grpc/context.rsmodel_gateway/src/routers/grpc/pipeline.rsmodel_gateway/src/routers/grpc/router.rsmodel_gateway/src/routers/http/mod.rsmodel_gateway/src/routers/http/pd_router.rsmodel_gateway/src/routers/http/router.rsmodel_gateway/src/routers/router_manager.rsmodel_gateway/src/worker/registry.rsmodel_gateway/tests/common/mock_worker.rsmodel_gateway/tests/routing/mod.rsmodel_gateway/tests/routing/model_alias_test.rsmodel_gateway/tests/routing/pd_routing_test.rs
2 participants
Summary
get_by_modelunderstands aliases/v1/responsesavailability gate, which runs before canonicalizationBehavior change
A request addressed to an alias is answered with the canonical model ID, matching how the OpenAI API reports the model it actually ran.
Tests
/v1/responsesgate: accepts aliases, still rejects unknown models and theunknownwildcardSummary by CodeRabbit
New Features
Bug Fixes
404 NOT_FOUNDresponses.Tests