Repository navigation
feat(gateway): API-key-aware /v1/models with upstream fan-out - #698
Conversation
github-actions
Bot
added
tests
📝 WalkthroughWalkthroughRouterManager centralizes /v1/models: it now accepts a reqwest::Client, extracts bearer tokens, fans out to upstream providers when a non-gateway key is present, parses upstream OpenAI/Anthropic JSON via a new parse_upstream, and falls back to local registry models. Per-router get_models handlers were removed. Changes
Sequence DiagramsequenceDiagram
participant Client
participant RouterManager
participant Auth
participant Upstream
participant WorkerRegistry
Client->>RouterManager: GET /v1/models (Authorization / x-api-key)
RouterManager->>Auth: extract bearer token / api-key
Auth-->>RouterManager: token
alt Token matches gateway_api_key
RouterManager->>WorkerRegistry: fetch local ModelCard list
WorkerRegistry-->>RouterManager: local models
RouterManager-->>Client: ListModelsResponse (registry models)
else Token present (provider key)
RouterManager->>Upstream: parallel GET /v1/models (with token)
alt First upstream returns 2xx
Upstream-->>RouterManager: JSON model list
RouterManager->>RouterManager: parse_upstream(json, provider)
RouterManager-->>Client: ListModelsResponse (upstream models)
else All upstreams fail
RouterManager->>WorkerRegistry: fetch local ModelCard list
WorkerRegistry-->>RouterManager: local models
RouterManager-->>Client: ListModelsResponse (local fallback)
end
else No token
RouterManager->>WorkerRegistry: fetch local ModelCard list
WorkerRegistry-->>RouterManager: local models
RouterManager-->>Client: ListModelsResponse (local models)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
|
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 |
mergify
Bot
added
the
needs-rebase
When a bearer token is present in a /v1/models request, the gateway now fans out to all healthy external upstreams using that token to discover available models (BYOK flow). If any upstream returns models, those are returned. If all upstreams fail or return nothing, the gateway falls back to self-hosted registry models. What changed: - model_gateway/src/routers/router_manager.rs: Rewrote get_models to extract bearer token and call fetch_upstream_models when present. Added fetch_upstream_models (deduplicates by URL, fans out in parallel) and fetch_models_from (per-upstream GET with apply_provider_headers). Added client field to RouterManager struct and updated constructor. - crates/protocols/src/models.rs: Added ListModelsResponse::parse_upstream that handles both OpenAI and Anthropic response schemas (both use data[].id) and converts to Vec<ModelCard> with provider inference. - Removed get_models overrides from all individual routers (http router, pd_router, openai router, anthropic router, gemini router) — the RouterManager is now the single implementation. - Deleted model_gateway/src/routers/anthropic/models.rs (only contained handle_list_models which is no longer needed). - Updated policy_registry_integration.rs test for new constructor sig. Why: Users sending their own provider API keys (e.g., OpenAI, Anthropic) should be able to discover which models are available upstream, rather than only seeing locally-registered models. This enables BYOK workflows. How: The /v1/models endpoint is a public route (no auth middleware). When a bearer token is present, RouterManager fans out GET /v1/models to each unique external worker URL in parallel, using apply_provider_headers for provider-specific auth formatting (Anthropic x-api-key vs standard Bearer). Response parsing lives in the protocols crate as ListModelsResponse::parse_upstream, keeping router_manager focused on orchestration. Individual router get_models implementations were removed to centralize the logic. Refs: #691 Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
6d13f31 to
984af9a
Compare
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 /v1/models API endpoint by enabling dynamic model discovery from external upstream providers when a user supplies their own API key. This change centralizes the model discovery logic within the RouterManager, eliminating redundant implementations across individual routers and improving the user experience for "bring your own key" (BYOK) scenarios by providing a comprehensive view of available models. 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
|
mergify
Bot
removed
the
needs-rebase
![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: 6d13f31807
ℹ️ 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".
![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 a significant feature for API-key-aware model discovery from upstream providers, centralizing /v1/models logic and implementing a robust fan-out mechanism with a flexible parser. However, it introduces a critical credential leakage vulnerability: the fan-out mechanism sends a user's API key for one provider to all other configured external providers. Additionally, there's a minor suggestion to improve the efficiency of collecting unique upstream URLs, while acknowledging the correct use of HashSet for uniqueness.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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 `@model_gateway/src/routers/router_manager.rs`:
- Around line 376-384: The current code maps unique_urls into futures and uses
futures::future::join_all then flattens all successful results, producing a
merged cross-provider list; change this to return the first successful provider
inventory instead: iterate over unique_urls in order and for each call and await
Self::fetch_models_from(self.client.clone(), base, auth.clone()) sequentially,
and as soon as a call returns a successful non-empty inventory (or otherwise
qualifies as a successful response per fetch_models_from's contract) return that
Vec immediately; remove the join_all + flatten logic so you do not merge results
across providers and preserve per-provider inventories (use the symbols
unique_urls and fetch_models_from to locate and update the code).
- Around line 481-487: The current fallback builds cards from
self.worker_registry.get_all() which includes external workers; change the
pipeline that collects registry models to exclude external workers by filtering
the iterator (e.g., use .get_all().iter().filter(|w| !w.is_external() ) or the
equivalent check on the worker kind) before flat_map(|w| w.models()) so only
self-hosted registry models are returned by the fallback.
- Around line 462-479: The current logic treats any Bearer token as a BYOK and
probes upstreams via fetch_upstream_models; add a short-circuit to skip upstream
fan-out when the token equals the gateway's configured key so
gateway-authenticated clients stay on the registry path. Concretely, before
calling fetch_upstream_models(bearer_token), check whether
bearer_token.as_deref() == Some(&self.gateway_key) (or the appropriate config
field name) and if so do not call fetch_upstream_models but continue to the
registry-handling path; otherwise proceed with the existing upstream probe and
ListModelsResponse flow.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 3141917a-4674-4b66-9400-1ded53f88706
📒 Files selected for processing (10)
crates/protocols/src/models.rsmodel_gateway/src/routers/anthropic/mod.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/routing/policy_registry_integration.rs
💤 Files with no reviewable changes (5)
- model_gateway/src/routers/openai/router.rs
- model_gateway/src/routers/http/pd_router.rs
- model_gateway/src/routers/anthropic/models.rs
- model_gateway/src/routers/anthropic/mod.rs
- model_gateway/src/routers/http/router.rs
- Fix test_v1_models and test_model_info_with_no_workers: wrap single router in RouterManager in test helper so get_models is handled by RouterManager (matching production behavior). - Fix test_health_generate_endpoint and test_get_server_info: delegate health_generate and get_server_info to inner router instead of returning stub responses, consistent with how get_model_info works. - Return first successful upstream inventory instead of merging all providers, preventing credential leakage to unrelated upstreams. - Parse Authorization header case-insensitively per RFC 7235. - Exclude external workers from registry fallback so only self-hosted models are returned when upstream discovery fails. Refs: #691 Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
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/router_manager.rs`:
- Around line 376-383: The code silently drops authentication when
HeaderValue::from_str(&format!("Bearer {bearer_token}")) fails, so change the
auth construction to handle the Result explicitly: call HeaderValue::from_str,
match Err to emit a warning (using the existing logger/tracing) that the bearer
token contained invalid characters (include the bearer_token or an obfuscated
version), and only pass Some(header) into Self::fetch_models_from when Ok; keep
passing auth.clone() as before but ensure auth is Option<HeaderValue> derived
from the matched result. Reference HeaderValue::from_str, the bearer_token
variable, the auth local, and Self::fetch_models_from in your change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 89ed7596-05fa-4271-bbf2-5f292cbcc089
📒 Files selected for processing (2)
model_gateway/src/routers/router_manager.rsmodel_gateway/tests/common/mod.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 37a7c507c7
ℹ️ 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".
| .headers() | ||
| .get(header::AUTHORIZATION) | ||
| .and_then(|h| h.to_str().ok()) |
There was a problem hiding this comment.
get_models only reads Authorization and derives a bearer token from it, so requests that use Anthropic-style credentials (x-api-key with anthropic-version) never trigger upstream discovery. Because the fallback path now excludes external workers, valid Anthropic BYOK callers can receive 503 No models available (or unrelated local models) in external-provider deployments even though their credentials are valid for /v1/messages.
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 the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/routers/router_manager.rs`:
- Around line 423-446: The fetch_models_from function lacks a per-request
timeout, so update the RequestBuilder produced by apply_provider_headers (the
req variable) to set a reasonable timeout (e.g., using
req.timeout(Duration::from_secs(5)) from reqwest) before calling .send(). Also
adjust the error branch that handles req.send().await to recognize and log
timeout errors distinctly (or include the error details) and return an empty
Vec; make these changes in fetch_models_from where req is defined and used.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: dc57168e-2580-4bcf-98c1-b5fd782b9dc7
📒 Files selected for processing (2)
model_gateway/src/routers/router_manager.rsmodel_gateway/tests/routing/test_openai_routing.rs
- Add gateway-key short-circuit: skip upstream fan-out when bearer token matches the configured gateway API key - Accept Anthropic-style x-api-key header in addition to Authorization Bearer for model discovery - Warn and return early when bearer token contains invalid header chars instead of silently proceeding with None auth - Switch from sequential upstream iteration to concurrent fan-out using futures::select_all, returning first successful non-empty result - Extract registry_models_response() helper to deduplicate fallback path - Update OpenAI routing tests to use RouterManager (get_models is now centralized) and add non-external worker for registry path test Refs: #691 Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
8ebe8e4 to
b4bba8e
Compare
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
There was a problem hiding this comment.
♻️ Duplicate comments (1)
model_gateway/src/routers/router_manager.rs (1)
430-444:⚠️ Potential issue | 🟠 MajorAdd a timeout to prevent hanging on slow or unresponsive upstreams.
The upstream request lacks an explicit timeout. If an upstream is slow or unresponsive, this can cause the request to hang, potentially affecting the overall
/v1/modelslatency since theselect_allloop will wait for all futures if none return non-empty results.🛡️ Proposed fix to add request timeout
async fn fetch_models_from( client: reqwest::Client, base_url: String, auth: Option<HeaderValue>, ) -> Vec<ModelCard> { let url = format!("{base_url}/v1/models"); - let req = apply_provider_headers(client.get(&url), &url, auth.as_ref()); + let req = apply_provider_headers(client.get(&url), &url, auth.as_ref()) + .timeout(std::time::Duration::from_secs(10)); let resp = match req.send().await { Ok(r) => r, Err(e) => { debug!("Failed to reach upstream {url}: {e}"); return Vec::new(); } };🤖 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 430 - 444, The fetch_models_from function currently awaits req.send().await with no timeout; update the request created by apply_provider_headers (the RequestBuilder returned for client.get(&url)) to set a per-request timeout (e.g., .timeout(Duration::from_secs(5))) before calling send(), import std::time::Duration, and handle the timeout error branch similarly to other errors (log a debug/error mentioning the url and that the request timed out and return Vec::new()). Ensure you modify the RequestBuilder chain produced by apply_provider_headers (used in fetch_models_from) rather than only the client-level config.
🤖 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/router_manager.rs`:
- Around line 430-444: The fetch_models_from function currently awaits
req.send().await with no timeout; update the request created by
apply_provider_headers (the RequestBuilder returned for client.get(&url)) to set
a per-request timeout (e.g., .timeout(Duration::from_secs(5))) before calling
send(), import std::time::Duration, and handle the timeout error branch
similarly to other errors (log a debug/error mentioning the url and that the
request timed out and return Vec::new()). Ensure you modify the RequestBuilder
chain produced by apply_provider_headers (used in fetch_models_from) rather than
only the client-level config.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: c9379099-2828-464d-9154-ca89502da1cd
📒 Files selected for processing (2)
model_gateway/src/routers/router_manager.rsmodel_gateway/tests/routing/test_openai_routing.rs
…oject#698) Signed-off-by: Simo Lin <linsimo.mark@gmail.com> Signed-off-by: Sydney Firmin <sydney.firmin@oracle.com>
2 participants
Summary
Implements API-key-aware
/v1/modelsfor BYOK (bring your own key) flows. When a bearer token is present, the gateway fans out to all healthy external upstreams to discover available models. Falls back to self-hosted registry models if upstreams fail or return nothing.Closes #691
What changed
model_gateway/src/routers/router_manager.rs: Rewroteget_modelsto extract bearer token and fan out to upstreams. Addedfetch_upstream_models(deduplicates by URL, parallel fan-out) andfetch_models_from(per-upstream GET withapply_provider_headers). Addedclient: reqwest::Clientfield.crates/protocols/src/models.rs: AddedListModelsResponse::parse_upstream— handles both OpenAI and Anthropic response schemas (data[].id) and converts toVec<ModelCard>with provider inference.get_modelsfrom all individual routers (http, pd, openai, anthropic, gemini) —RouterManageris now the single implementation.model_gateway/src/routers/anthropic/models.rs— only containedhandle_list_models, no longer needed.policy_registry_integration.rsfor newRouterManager::newsignature.Why
Users sending their own provider API keys should discover which models are available upstream, not just locally-registered ones. The previous implementation had each router duplicating
/v1/modelslogic with no upstream discovery.How
/v1/modelsis a public route (no auth middleware). When a bearer token is present:RouterManager::fetch_upstream_modelsgets healthy external workers, deduplicates by URLGET /v1/modelsto each unique upstream in parallelapply_provider_headershandles provider-specific auth (Anthropicx-api-keyvs standardBearer)ListModelsResponse::parse_upstream(protocols crate) parses both OpenAI and Anthropic JSON schemasTest plan
cargo build -p smg— compiles cleanlycargo clippy -p smg --all-targets --all-features -- -D warnings— no warningscargo test -p smg— existing tests passcurl /v1/models(no auth) → registry modelscurl -H "Authorization: Bearer <provider-key>" /v1/models→ upstream models (if external workers configured)curl -H "Authorization: Bearer <invalid-key>" /v1/models→ falls back to registry modelsSummary by CodeRabbit
New Features
Bug Fixes & Improvements