refactor(gateway): unify /v1/models response across all routers - #692
Conversation
📝 WalkthroughWalkthroughAdds a public Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Router as HTTP Router
participant WRegistry as Worker Registry
participant Workers as Workers
rect rgba(100, 150, 200, 0.5)
Note over Client,Workers: Before (Proxy Approach)
Client->>Router: GET /v1/models
Router->>Workers: Forward request to one worker
Workers->>Router: Proxy response
Router->>Client: Return proxied response
end
rect rgba(150, 200, 100, 0.5)
Note over Client,Workers: After (Local Aggregation)
Client->>Router: GET /v1/models
Router->>WRegistry: Request worker list
WRegistry->>Router: Return workers
Router->>Workers: Query each worker for model cards
Workers->>Router: Return model cards
Router->>Router: Aggregate, dedupe, build ListModelsResponse
Router->>Client: Return aggregated JSON response
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related issues
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
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 standardizes the Highlights
Changelog
Activity
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. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
|
Hi @slin1237, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
88a0152 to
fddcd04
Compare
There was a problem hiding this comment.
Code Review
This pull request is a great step towards unifying the /v1/models endpoint, which will improve client reliability. The introduction of shared ModelObject and ListModelsResponse types and the refactoring of routers to use them is a solid approach. However, I've identified a recurring critical issue across multiple files where the new from_model_cards function is called with an incorrect iterator type, which will lead to compilation errors. I've provided suggestions to fix these issues and, in some cases, improve performance by avoiding unnecessary memory allocations.
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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/gemini/router.rs`:
- Around line 96-105: The get_models method currently always returns a 200
ListModelsResponse even when no model cards exist; update async fn get_models to
check the collected cards (variable cards) and if cards.is_empty() return a 503
Service Unavailable response (matching RouterManager::get_models fallback)
instead of constructing ListModelsResponse; locate get_models in router.rs and
return an appropriate error Response (StatusCode::SERVICE_UNAVAILABLE) when
cards are empty so the manager can fall back consistently.
In `@model_gateway/src/routers/http/pd_router.rs`:
- Around line 1241-1249: The get_models handler currently builds cards from
self.worker_registry.get_all().iter().flat_map(|w| w.models()) and always
returns 200 with a serialized ListModelsResponse; change it to check if the
collected cards vector is empty and, if so, return a 503
(StatusCode::SERVICE_UNAVAILABLE) response before constructing
openai_protocol::models::ListModelsResponse so PD mode matches the RouterManager
fallback behavior; locate the async fn get_models, add a cards.is_empty() branch
that returns a SERVICE_UNAVAILABLE response (matching the manager fallback) and
only serialize and return the 200 JSON when cards is non-empty.
In `@model_gateway/src/routers/http/router.rs`:
- Around line 675-682: The /v1/models listing currently collects model cards
from self.worker_registry.get_all(), which can include workers this router
cannot route to; update the collection to only include workers this router can
serve by filtering the registry for WorkerType::Regular and ConnectionMode::Http
(the same criteria used in select_worker_for_model) before calling .models() and
passing the cards to
openai_protocol::models::ListModelsResponse::from_model_cards so advertised
models actually resolve via this router.
In `@model_gateway/src/routers/openai/router.rs`:
- Around line 555-558: The ListModelsResponse building is losing provider
metadata because refresh_worker_models recreates ModelCard with
ModelCard::new(id) and drops provider info; update refresh_worker_models (or the
code that repopulates external_workers) to preserve the original
provider/provider-specific metadata when reconstructing ModelCard instances
(e.g., copy provider, owner, and any upstream tags into the new ModelCard or
construct from the existing card instead of ModelCard::new), so that
ListModelsResponse::from_model_cards(external_workers.iter().flat_map(|w|
w.models())) sees correct provider/owned_by values for upstream catalogs.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: dc2782d6-bac9-43f0-98e5-22003554a0d2
📒 Files selected for processing (11)
clients/openapi-gen/src/main.rscrates/protocols/src/lib.rscrates/protocols/src/model_card.rscrates/protocols/src/models.rsmodel_gateway/src/routers/anthropic/models.rsmodel_gateway/src/routers/gemini/router.rsmodel_gateway/src/routers/http/pd_router.rsmodel_gateway/src/routers/http/router.rsmodel_gateway/src/routers/openai/router.rsmodel_gateway/src/routers/router_manager.rsmodel_gateway/tests/api/api_endpoints_test.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 88a0152920
ℹ️ 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".
fddcd04 to
944a0fe
Compare
There was a problem hiding this comment.
♻️ Duplicate comments (2)
model_gateway/src/routers/http/router.rs (1)
674-682:⚠️ Potential issue | 🟠 MajorRestrict
/v1/modelsto models this router can actually serve.
select_worker_for_model()only routes toWorkerType::Regular+ConnectionMode::Http, but this listing pulls cards from every registered worker. That lets PD/gRPC/external-only models show up here and then fail on the next request through this router. Filter with the same criteria before serializing, and return503when nothing matches so the behavior stays aligned with the manager fallback.Suggested fix
async fn get_models(&self, _req: Request<Body>) -> Response { let cards = self .worker_registry - .get_all() - .iter() + .get_workers_filtered( + None, + Some(WorkerType::Regular), + Some(ConnectionMode::Http), + None, + false, + ) + .into_iter() .flat_map(|w| w.models()) .collect::<Vec<_>>(); + + if cards.is_empty() { + return error::service_unavailable( + "no_available_workers", + "No regular HTTP models available", + ); + } + let resp = openai_protocol::models::ListModelsResponse::from_model_cards(cards); (StatusCode::OK, Json(resp)).into_response() }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/http/router.rs` around lines 674 - 682, get_models currently gathers model cards from every registered worker but select_worker_for_model only serves WorkerType::Regular with ConnectionMode::Http, so filter the worker list the same way before collecting cards: call worker_registry.get_all(), filter workers by WorkerType::Regular and connection_mode == ConnectionMode::Http (matching select_worker_for_model's logic), then flat_map their models() and pass those cards to ListModelsResponse::from_model_cards; if the filtered list yields zero cards, return a 503 (Service Unavailable) response instead of 200 to match manager fallback behavior.model_gateway/src/routers/http/pd_router.rs (1)
1241-1249:⚠️ Potential issue | 🟠 MajorOnly list models backed by a complete PD pair.
This gathers cards from every registered worker, but
select_pd_pair()later succeeds only when the samemodel_idhas both a prefill and a decode worker. As a result,/v1/modelscan advertise regular-worker models or half-configured PD models that this router will immediately reject on the next request. It also still returns200with an empty list when no routable pair exists. Build the response from the intersection of prefill/decode model IDs and return503when that set is empty.Suggested fix
async fn get_models(&self, _req: Request<Body>) -> Response { - let cards = self - .worker_registry - .get_all() - .iter() - .flat_map(|w| w.models()) - .collect::<Vec<_>>(); + let decode_model_ids = self + .worker_registry + .get_decode_workers() + .into_iter() + .flat_map(|w| w.models().into_iter().map(|card| card.id)) + .collect::<std::collections::HashSet<_>>(); + + let cards = self + .worker_registry + .get_prefill_workers() + .into_iter() + .flat_map(|w| w.models()) + .filter(|card| decode_model_ids.contains(&card.id)) + .collect::<Vec<_>>(); + + if cards.is_empty() { + return error::service_unavailable( + "no_available_workers", + "No routable PD models available", + ); + } + let resp = openai_protocol::models::ListModelsResponse::from_model_cards(cards); (StatusCode::OK, axum::Json(resp)).into_response() }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/http/pd_router.rs` around lines 1241 - 1249, get_models currently collects cards from all workers but must only advertise models that have both a prefill and decode PD pair; modify get_models to query worker_registry for prefill and decode workers, build sets of model_ids from each (using the same methods that select_pd_pair relies on), compute the intersection, filter the collected model cards to only those whose model_id is in that intersection (use openai_protocol::models::ListModelsResponse::from_model_cards as before), and if the resulting set is empty return a 503 response instead of 200; reference symbols: get_models, worker_registry, select_pd_pair, ListModelsResponse, models().
🤖 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/routers/http/pd_router.rs`:
- Around line 1241-1249: get_models currently collects cards from all workers
but must only advertise models that have both a prefill and decode PD pair;
modify get_models to query worker_registry for prefill and decode workers, build
sets of model_ids from each (using the same methods that select_pd_pair relies
on), compute the intersection, filter the collected model cards to only those
whose model_id is in that intersection (use
openai_protocol::models::ListModelsResponse::from_model_cards as before), and if
the resulting set is empty return a 503 response instead of 200; reference
symbols: get_models, worker_registry, select_pd_pair, ListModelsResponse,
models().
In `@model_gateway/src/routers/http/router.rs`:
- Around line 674-682: get_models currently gathers model cards from every
registered worker but select_worker_for_model only serves WorkerType::Regular
with ConnectionMode::Http, so filter the worker list the same way before
collecting cards: call worker_registry.get_all(), filter workers by
WorkerType::Regular and connection_mode == ConnectionMode::Http (matching
select_worker_for_model's logic), then flat_map their models() and pass those
cards to ListModelsResponse::from_model_cards; if the filtered list yields zero
cards, return a 503 (Service Unavailable) response instead of 200 to match
manager fallback behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: cd8fc444-7697-4b04-b793-7d8f6efe9add
📒 Files selected for processing (12)
clients/openapi-gen/src/main.rscrates/protocols/src/lib.rscrates/protocols/src/model_card.rscrates/protocols/src/models.rsmodel_gateway/src/routers/anthropic/models.rsmodel_gateway/src/routers/anthropic/router.rsmodel_gateway/src/routers/gemini/router.rsmodel_gateway/src/routers/http/pd_router.rsmodel_gateway/src/routers/http/router.rsmodel_gateway/src/routers/openai/router.rsmodel_gateway/src/routers/router_manager.rsmodel_gateway/tests/api/api_endpoints_test.rs
All routers now return the same OpenAI-format response built from in-memory ModelCard data instead of ad-hoc JSON or upstream proxying. What changed: - crates/protocols/src/models.rs: new module with ModelObject and ListModelsResponse types, including from_model_cards() constructor that deduplicates by model ID - crates/protocols/src/model_card.rs: add owned_by() and into_model_object() methods for ModelCard-to-response conversion - crates/protocols/src/lib.rs: register new models module - model_gateway/src/routers/openai/router.rs: replace ad-hoc json!() with ListModelsResponse::from_model_cards(), drop non-standard aliases/model_type/primary_model fields - model_gateway/src/routers/anthropic/models.rs: replace HTTP proxy approach with building from filtered Anthropic ModelCards, removing ~50 lines of proxy/size-check logic - model_gateway/src/routers/http/router.rs: replace proxy_get_request with worker registry ModelCard-based response - model_gateway/src/routers/http/pd_router.rs: replace proxy_to_first_prefill_worker with registry-based response - model_gateway/src/routers/gemini/router.rs: add get_models() override (previously fell through to 501 NOT_IMPLEMENTED) - model_gateway/src/routers/router_manager.rs: replace fallback ad-hoc JSON with ListModelsResponse, remove unused serde_json::Value - clients/openapi-gen/src/main.rs: swap to new models::ListModelsResponse - model_gateway/tests/api/api_endpoints_test.rs: update test to match registry-based model ID and owned_by values Why: Each router had its own /v1/models implementation with inconsistent formats: OpenAI router included non-standard fields (aliases, model_type, primary_model), Anthropic router proxied raw upstream responses, HTTP/PD routers proxied to backends, Gemini returned 501, and the router manager fallback used owned_by: "local". This made client integration unreliable across router types. How: Introduced a shared models module in the protocols crate with proper serde types. ModelCard gains owned_by() (None -> "self_hosted", provider -> provider.as_str()) and into_model_object() (consuming conversion to avoid cloning). All routers now collect ModelCards from their worker registries and delegate to ListModelsResponse::from_model_cards(). Refs: #691 Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
944a0fe to
1830a36
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1830a3690e
ℹ️ 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".
… workers Address PR review feedback: - model_gateway/src/routers/gemini/router.rs: return 503 SERVICE_UNAVAILABLE when worker registry has no model cards, matching router_manager fallback - model_gateway/src/routers/http/pd_router.rs: same empty-cards → 503 check - model_gateway/src/routers/http/router.rs: filter to Regular+Http workers (matching select_worker_for_model criteria) so /v1/models only advertises models this router can actually route to; add empty-cards → 503 check Why: without these checks, individual routers returned 200 with an empty list when no models exist, while the router_manager fallback returned 503. The HTTP router also advertised models from worker types it can't route to. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
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/routers/router_manager.rs (1)
379-401:⚠️ Potential issue | 🟠 MajorMake the default
/v1/modelspath deterministic instead of delegating to an arbitrary router.For requests without a provider-specific header,
select_router_for_request(..., None)gives every non-PD router the same score and keeps the first entry fromrouters_snapshot. After this PR those routers no longer return the same dataset: Anthropic is provider-scoped, OpenAI refreshes external cards before listing, and the regular/PD routers dump the registry directly. The manager can now return different/v1/modelspayloads for the same request depending on snapshot iteration order.Please make the manager own the default unified response path, and reserve router delegation for explicitly provider-scoped requests.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/router_manager.rs` around lines 379 - 401, get_models currently delegates to select_router_for_request when no provider header is present, which makes the default /v1/models response nondeterministic; change it so the manager itself returns the unified default list and only delegates to a router when the request is explicitly provider-scoped (e.g., contains "anthropic-version" or other provider-specific header/param). Specifically, in get_models replace the select_router_for_request(...) fallback with manager-owned handling: when parts.headers contains "anthropic-version" (or other explicit provider header) find that router via self.routers.get(&router_ids::HTTP_ANTHROPIC) and forward the reconstructed Request (Request::from_parts(parts, body)) to router.get_models(...); otherwise, do not call select_router_for_request and instead call the manager's internal unified listing logic (the code that builds the default /v1/models payload) and return that Response directly. Ensure you keep the existing special-case check for response.status() != StatusCode::NOT_IMPLEMENTED when forwarding to a router.
🤖 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/http/router.rs`:
- Around line 674-682: The get_models handler is returning stale external models
because it calls worker_registry.get_all() and then w.models() without
performing the provider refresh that openai/router.rs does; update get_models to
either filter out external workers (e.g., skip workers where
Worker::is_external() is true) or invoke the provider refresh path before
collecting models (call the equivalent refresh method used in
model_gateway/src/routers/openai/router.rs for each external worker or use a
registry-level refresh method) so the aggregate built from
worker_registry.get_all().iter().flat_map(|w| w.models()) returns up-to-date
model data.
---
Outside diff comments:
In `@model_gateway/src/routers/router_manager.rs`:
- Around line 379-401: get_models currently delegates to
select_router_for_request when no provider header is present, which makes the
default /v1/models response nondeterministic; change it so the manager itself
returns the unified default list and only delegates to a router when the request
is explicitly provider-scoped (e.g., contains "anthropic-version" or other
provider-specific header/param). Specifically, in get_models replace the
select_router_for_request(...) fallback with manager-owned handling: when
parts.headers contains "anthropic-version" (or other explicit provider header)
find that router via self.routers.get(&router_ids::HTTP_ANTHROPIC) and forward
the reconstructed Request (Request::from_parts(parts, body)) to
router.get_models(...); otherwise, do not call select_router_for_request and
instead call the manager's internal unified listing logic (the code that builds
the default /v1/models payload) and return that Response directly. Ensure you
keep the existing special-case check for response.status() !=
StatusCode::NOT_IMPLEMENTED when forwarding to a router.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 1fb611c7-8e7b-42d8-a3bb-6cac9f17dc9e
📒 Files selected for processing (12)
clients/openapi-gen/src/main.rscrates/protocols/src/lib.rscrates/protocols/src/model_card.rscrates/protocols/src/models.rsmodel_gateway/src/routers/anthropic/models.rsmodel_gateway/src/routers/anthropic/router.rsmodel_gateway/src/routers/gemini/router.rsmodel_gateway/src/routers/http/pd_router.rsmodel_gateway/src/routers/http/router.rsmodel_gateway/src/routers/openai/router.rsmodel_gateway/src/routers/router_manager.rsmodel_gateway/tests/api/api_endpoints_test.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 45cfbf7331
ℹ️ 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
🤖 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/http/pd_router.rs`:
- Around line 1241-1252: get_models currently lists every model from all workers
but must only advertise models that exist on both sides of a PD pair; change
get_models to compute the set of model ids available on prefill workers and the
set available on decode workers (using worker_registry.get_all() and each
Worker.models()), intersect those sets, then build the ListModelsResponse only
from model cards whose model id is in that intersection so select_pd_pair() will
succeed for advertised models.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 0bb32975-8c1c-4f11-8d42-1b9b39264508
📒 Files selected for processing (3)
model_gateway/src/routers/gemini/router.rsmodel_gateway/src/routers/http/pd_router.rsmodel_gateway/src/routers/http/router.rs
Summary
/v1/modelsto return the same OpenAI-format response across all 6 router typesModelObject/ListModelsResponsetypes in the protocols crateModelCarddata instead of ad-hoc JSON or upstream proxyingRefs: #691
What changed
crates/protocols/src/models.rsModelObject,ListModelsResponse,from_model_cards()with dedupcrates/protocols/src/model_card.rsowned_by()andinto_model_object()crates/protocols/src/lib.rsmodelsmodulemodel_gateway/src/routers/openai/router.rsjson!()withListModelsResponse, drop non-standardaliases/model_type/primary_modelfieldsmodel_gateway/src/routers/anthropic/models.rsmodel_gateway/src/routers/http/router.rsproxy_get_requestwith registry-based responsemodel_gateway/src/routers/http/pd_router.rsproxy_to_first_prefill_workerwith registry-based responsemodel_gateway/src/routers/gemini/router.rsget_models()override (was returning 501)model_gateway/src/routers/router_manager.rsListModelsResponseclients/openapi-gen/src/main.rsmodels::ListModelsResponsefor OpenAPI schemamodel_gateway/tests/api/api_endpoints_test.rsWhy
Each router had its own
/v1/modelsimplementation with inconsistent formats:json!()with non-standard fields (aliases,model_type,primary_model)json!()withowned_by: "local"This made client integration unreliable across router types.
How
Introduced a
modelsmodule in the protocols crate with proper serde types.ModelCardgainsowned_by()(None→"self_hosted", provider →provider.as_str()) andinto_model_object()(consuming conversion). All routers collectModelCards from their worker registries and delegate toListModelsResponse::from_model_cards().Response format:
{ "object": "list", "data": [ { "id": "model-id", "object": "model", "created": 0, "owned_by": "self_hosted" } ] }Test plan
cargo build— all crates compilecargo test -p openai-protocol— 4 new unit tests pass (serialization, roundtrip deserialization, dedup, owned_by mapping)cargo test -p smg— all 93 integration tests passtest_v1_modelsintegration test to match new registry-based model ID andowned_byvaluesSummary by CodeRabbit
New Features
Bug Fixes
Tests