Repository navigation
feat(interactions): Register Gemini Router - #679
Merged
Merged
Conversation
XinyueZhang369
requested review from
CatherineSue,
key4ng and
slin1237
as code owners
March 9, 2026 16:44
github-actions
Bot
added
tests
📝 WalkthroughWalkthroughThis pull request adds Gemini backend routing support to the model gateway. It introduces a new Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 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)
Comment |
Contributor
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 model gateway by integrating Google's Gemini API. It establishes a new routing mode that allows the gateway to process non-streaming, non-tool interaction requests for Gemini models. The changes span configuration, worker management, and a dedicated request-response flow, ensuring seamless communication with Gemini-compatible endpoints and laying the groundwork for future streaming and tool-use capabilities. Highlights
Changelog
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
|
![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: 64f2a656bf
ℹ️ 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}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces support for Google's Gemini API as a new routing mode in the model gateway. A high-severity security vulnerability has been identified: the gateway leaks user-supplied provider-specific API keys (e.g., Gemini keys) to other configured external providers (e.g., OpenAI, Anthropic) during model discovery. While the issue is identified here, addressing such cross-cutting security concerns, especially those requiring design decisions, should ideally be handled in a dedicated pull request. Additionally, I've identified several opportunities for improvement related to code duplication, robust error handling (preferring explicit logging over silent failures), and performance optimizations, which are detailed in the specific comments.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 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/job_queue.rs`:
- Around line 561-595: The Gemini branch currently builds a generic external
worker config and loses the RoutingMode::Gemini identity, causing discovery
(discover_models.rs) to re-infer provider from config.url; fix by carrying the
routing mode/provider through the worker config/workflow metadata instead of
relying on URL inference: update build_external_worker_config (and its callers
in this branch) to accept and set an explicit provider/routing_mode field
(preserving RoutingMode::Gemini), ensure Job::AddWorker's config includes that
provider flag, and update discover_models.rs to prefer the explicit
provider/routing_mode on the worker config/workflow before falling back to
URL-based heuristics so Gemini proxies/custom hostnames correctly use Gemini
discovery and x-goog-api-key paths.
In `@model_gateway/src/routers/gemini/steps/worker_selection.rs`:
- Around line 154-160: The fragile URL-based check using
worker.url().contains("googleapis.com") should be replaced with the
provider-type check used elsewhere: test for ProviderType::Gemini (e.g.
worker.provider_type() == ProviderType::Gemini or matching worker.provider())
and set models_path to "/v1beta/models" for Gemini and "/v1/models" otherwise
(update the models_path assignment where worker.url() is currently used); also
add the necessary import/qualifier for ProviderType to mirror the logic in
discover_models.rs so proxied/custom domains are handled consistently.
In `@model_gateway/src/routers/gemini/utils.rs`:
- Around line 15-18: The code currently sets user_auth from headers.and_then(|h|
h.get("x-goog-api-key").cloned()) which treats an empty header value as present
and prevents fallback to worker_api_key; change the logic so user_auth only wins
if the header value is non-empty (trimmed), e.g. extract the header string,
check that it's not empty after trimming, and only then convert to HeaderValue;
otherwise let worker_api_key be used via the existing
worker_api_key.and_then(|k| HeaderValue::from_str(k).ok()) fallback. Ensure you
update the expression that builds user_auth (and the final
user_auth.or_else(...) chain) to perform the emptiness check before constructing
the HeaderValue.
In `@model_gateway/src/routers/header_utils.rs`:
- Around line 177-183: The Gemini branch currently calls auth.to_str() and uses
strip_prefix("Bearer ") which is case-sensitive; update the handling in the
ApiProvider::Gemini block so you detect and remove a "Bearer " scheme
case-insensitively before setting x-goog-api-key: e.g. get the auth_str from
auth.to_str(), check auth_str.to_ascii_lowercase().starts_with("bearer ") (or
compare a lowercased slice) and if true slice off the first 7 bytes to produce
api_key (also trim surrounding whitespace), otherwise use auth_str as-is, then
pass that api_key to req.header("x-goog-api-key", api_key) to ensure lowercase
"bearer" is handled correctly.
In `@model_gateway/src/routers/router_manager.rs`:
- Around line 609-631: route_interactions currently selects the model directly
from model_id or body fields and calls select_router_for_request, skipping
IGW-specific resolution; update route_interactions to call resolve_model_id (the
same helper used by other routes) with headers, model_id, and
body.model/body.agent to derive the final selected_model before calling
select_router_for_request, so IGW fallback/validation is applied consistently
and error handling matches other route methods (refer to resolve_model_id,
route_interactions, select_router_for_request, and InteractionsRequest to locate
the relevant code).
In `@model_gateway/tests/api/interactions_api_test.rs`:
- Around line 145-171: The test_interactions_model_not_found currently accepts
503 which masks discovery/refresh failures; change the assertion in
test_interactions_model_not_found to require StatusCode::NOT_FOUND only (remove
the || StatusCode::SERVICE_UNAVAILABLE branch) so a missing-model request routed
via RouterFactory::create_gemini_router and router.route_interactions(...) must
produce 404; if you need to cover refresh/discovery failures, add a separate
test that simulates an unhealthy upstream and asserts SERVICE_UNAVAILABLE.
- Around line 279-309: The test test_interactions_forwards_api_key_header
currently passes because MockGeminiServer::check_auth accepts either
x-goog-api-key or Authorization, so add an explicit assertion that the mock
actually received the Gemini-specific header: update the mock (or the test) to
capture the incoming header name/value in MockGeminiServer::check_auth (or via
new_with_auth) and assert it equals "x-goog-api-key: test-key-123" in
test_interactions_forwards_api_key_header, or alternatively change
MockGeminiServer::check_auth to reject Authorization headers and only accept
x-goog-api-key so the test fails if bearer-only forwarding occurs.
In `@model_gateway/tests/common/mock_gemini_server.rs`:
- Around line 181-207: The mock_models function currently returns an
OpenAI-style payload; update it to validate auth by calling check_auth(State)
and to emit Gemini's payload shape: a JSON object with "models": [ { "name":
"models/..." }, ... ]. Keep the same model identifiers from mock_models
(gemini-2.5-flash, gemini-2.5-pro, deep-research-pro-preview-12-2025) but wrap
them as "models/{id}" in the "name" fields, and return that JSON via
Json(...).into_response() so the mock matches the production discovery contract
checked in discover_models.rs.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 19a72d8e-4ee4-496d-9dfb-17c6e37780a6
📒 Files selected for processing (25)
crates/protocols/src/interactions.rsmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/config/validation.rsmodel_gateway/src/core/job_queue.rsmodel_gateway/src/core/steps/worker/external/discover_models.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/factory.rsmodel_gateway/src/routers/gemini/context.rsmodel_gateway/src/routers/gemini/mod.rsmodel_gateway/src/routers/gemini/router.rsmodel_gateway/src/routers/gemini/state.rsmodel_gateway/src/routers/gemini/steps/non_stream_execution.rsmodel_gateway/src/routers/gemini/steps/request_building.rsmodel_gateway/src/routers/gemini/steps/response_processing.rsmodel_gateway/src/routers/gemini/steps/worker_selection.rsmodel_gateway/src/routers/gemini/utils.rsmodel_gateway/src/routers/header_utils.rsmodel_gateway/src/routers/mod.rsmodel_gateway/src/routers/router_manager.rsmodel_gateway/src/server.rsmodel_gateway/tests/api/interactions_api_test.rsmodel_gateway/tests/api/mod.rsmodel_gateway/tests/common/mock_gemini_server.rsmodel_gateway/tests/common/mod.rs
💤 Files with no reviewable changes (1)
- model_gateway/src/routers/gemini/state.rs
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
There was a problem hiding this comment.
Actionable comments posted: 9
🤖 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/external/discover_models.rs`:
- Around line 232-245: The discovery currently switches models_path and auth
header based on provider derived from the URL (ProviderType::from_url in
execute()), which fails for Gemini behind proxies; instead drive the branch from
the configured routing/provider metadata (e.g., use the configured provider
field such as config.provider or route.provider) rather than the URL-derived
provider. Update the code around models_path/models_url and header selection
(the provider, models_path, models_url, and the req.header(...) branches) to
consult the configured provider metadata and use that to choose "/v1beta/models"
and "x-goog-api-key" for Gemini, leaving other providers to the default
"/v1/models" and Bearer header.
In `@model_gateway/src/routers/gemini/steps/non_stream_execution.rs`:
- Around line 87-105: The code currently calls
worker.circuit_breaker().record_failure() for any non-success HTTP status and
returns internal_error on JSON parse failures; update the logic so that
circuit_breaker.record_failure() is only invoked for upstream 5xx (server)
responses (i.e., check response.status().is_server_error()), while 4xx responses
should not trip the breaker and should return the upstream status and sanitized
body as currently done; also change the JSON parse error branch (where
response.json().await fails and error::internal_error is used) to return a 502
Bad Gateway style error (use the project's bad-gateway helper or return a 502
status with a descriptive message) since parse failures are an upstream contract
error rather than an internal error.
In `@model_gateway/src/routers/gemini/steps/request_building.rs`:
- Around line 83-86: The code in request_building.rs currently inserts "model"
into the serialized request when ctx.input.model_id is Some (see
ctx.input.model_id and obj.insert("model"...)) but does not remove or reject an
existing "agent" field, which can cause a conflicting payload; update the logic
in the request-building block (where obj is modified) to either reject when both
ctx.input.model_id and an "agent" key are present (return an error early) or
normalize by removing obj.remove("agent") before inserting "model" so the path
model wins and only "model" is forwarded.
In `@model_gateway/src/routers/gemini/steps/response_processing.rs`:
- Around line 67-68: The code unconditionally echoes the caller's store flag via
obj.insert("store".to_string(), Value::Bool(req.store)); which advertises
durability the gateway doesn't provide yet; change the logic that sets "store"
so it only returns true when persistence is actually implemented (or when a
runtime/feature flag indicates persistence exists). Concretely, replace the
unconditional insert with a guarded value (e.g. Value::Bool(persistence_enabled
&& req.store) or always Value::Bool(false) until Phase 5 is implemented) in the
same response construction location (the obj.insert call) so responses never
claim storage unless real persistence is present.
In `@model_gateway/src/routers/gemini/steps/worker_selection.rs`:
- Around line 154-160: The current URL substring check (is_gemini =
worker.url().contains("googleapis.com")) is brittle; update the detection logic
used to set is_gemini and models_path to rely on explicit provider/metadata
fields instead of URL matching: read a dedicated field such as worker.provider()
or worker.metadata().get("provider") and treat values like "google" or "gemini"
as Gemini providers, set models_path accordingly ("/v1beta/models" for Gemini
else "/v1/models"), and only fall back to a URL contains check as a last-resort
safeguard; change references to is_gemini and models_path in this file
(worker.url(), is_gemini, models_path) to use the new provider/metadata-based
check.
- Around line 150-165: The HTTP call in refresh_worker_models currently uses
client.get(&url) without an explicit per-request timeout; update the request
built in refresh_worker_models (the req variable created from client.get(&url)
and passed to apply_provider_headers) to include a reasonable
.timeout(Duration::from_secs(...)) so the discovery request fails fast for
slow/unresponsive workers (import std::time::Duration and choose an appropriate
timeout, e.g., 2–5s). Ensure the timeout is applied before sending the request
so apply_provider_headers and subsequent await use the timed request.
In `@model_gateway/src/routers/gemini/utils.rs`:
- Around line 10-18: The extract_gemini_auth_header helper currently only checks
x-goog-api-key and the worker_api_key; update it to mirror header_utils.rs by
also accepting an Authorization: Bearer <token> header as a valid Gemini API key
source. Implement precedence: prefer headers.get("x-goog-api-key") first, then
parse headers.get("authorization") for a Bearer token (strip the "Bearer "
prefix and validate), and finally fall back to worker_api_key; return a
HeaderValue built from the selected token (using
HeaderValue::from_str(...).ok()). Ensure you reference the
extract_gemini_auth_header function and HeaderMap/worker_api_key parameters when
making the change.
In `@model_gateway/tests/common/mock_gemini_server.rs`:
- Around line 181-207: The mock_models handler returns an OpenAI-style schema
and doesn't enforce Gemini auth; update the mock_models function (and its use of
State<Arc<MockServerState>>) to return Gemini's contract: a JSON object with
"models": [{ "name": "models/<model-id>" , ... }] instead of "data", and
require/check the Gemini auth header on requests to /v1beta/models (validate the
Authorization header or call into MockServerState to assert the expected bearer
token) so tests exercise the new parser/header behavior.
In `@model_gateway/tests/common/mod.rs`:
- Around line 412-429: create_test_context() seeds Gemini external workers but
create_test_context_with_parsers() and create_test_context_with_mcp_config() do
not; update those helper functions to mirror the Gemini worker registration
block from the diff: detect RoutingMode::Gemini { worker_urls, .. } and for each
url build the same models vector and Arc<dyn Worker> via
BasicWorkerBuilder::new(url).worker_type(WorkerType::Regular).runtime_type(RuntimeType::External).models(models).build()
then call app_context.worker_registry.register(worker) so the registry is
populated for Gemini in all test-context helpers.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: ef83b86c-8c3f-42a5-9aac-cb8f1e7a80bf
📒 Files selected for processing (25)
crates/protocols/src/interactions.rsmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/config/validation.rsmodel_gateway/src/core/job_queue.rsmodel_gateway/src/core/steps/worker/external/discover_models.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/factory.rsmodel_gateway/src/routers/gemini/context.rsmodel_gateway/src/routers/gemini/mod.rsmodel_gateway/src/routers/gemini/router.rsmodel_gateway/src/routers/gemini/state.rsmodel_gateway/src/routers/gemini/steps/non_stream_execution.rsmodel_gateway/src/routers/gemini/steps/request_building.rsmodel_gateway/src/routers/gemini/steps/response_processing.rsmodel_gateway/src/routers/gemini/steps/worker_selection.rsmodel_gateway/src/routers/gemini/utils.rsmodel_gateway/src/routers/header_utils.rsmodel_gateway/src/routers/mod.rsmodel_gateway/src/routers/router_manager.rsmodel_gateway/src/server.rsmodel_gateway/tests/api/interactions_api_test.rsmodel_gateway/tests/api/mod.rsmodel_gateway/tests/common/mock_gemini_server.rsmodel_gateway/tests/common/mod.rs
💤 Files with no reviewable changes (1)
- model_gateway/src/routers/gemini/state.rs
CatherineSue
reviewed
Mar 9, 2026
Member
There was a problem hiding this comment.
Thanks for the work on this! Some feedback below.
Size: At 1380+ additions / 25 files, this is above our 1k line guideline. Consider splitting model discovery (GeminiModelsResponse, worker_selection refresh) into a separate PR — perhaps focus on /v1/models support first, then the interactions router steps.
Reference: Please add a reference to PR #417 in the description since this builds directly on that scaffolding.
Lint: Looks like CI lint is currently failing — please fix.
See inline comments for specifics.
CatherineSue
reviewed
Mar 9, 2026
CatherineSue
reviewed
Mar 9, 2026
XinyueZhang369
force-pushed
the
xz/non-stream-interactions-request
branch
from
March 9, 2026 17:42
64f2a65 to
7784367
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 77843671cd
ℹ️ 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: 2
♻️ Duplicate comments (12)
model_gateway/tests/common/mod.rs (1)
412-429:⚠️ Potential issue | 🟠 MajorMirror this Gemini worker registration in the other test-context helpers.
create_test_context()now seeds Gemini workers, butcreate_test_context_with_parsers()andcreate_test_context_with_mcp_config()still only register OpenAI workers. Tests using either helper withRoutingMode::Geminiwill start with an empty registry.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/tests/common/mod.rs` around lines 412 - 429, create_test_context() seeds Gemini external workers but create_test_context_with_parsers() and create_test_context_with_mcp_config() still only register OpenAI workers; update both helpers to mirror the Gemini worker registration logic: detect RoutingMode::Gemini and for each url build the same Vec<ModelCard> and Arc<dyn Worker> via BasicWorkerBuilder::new(url).worker_type(WorkerType::Regular).runtime_type(RuntimeType::External).models(models).build() and call app_context.worker_registry.register(worker) so the registry is populated the same way as in create_test_context().model_gateway/src/routers/gemini/steps/non_stream_execution.rs (2)
87-105:⚠️ Potential issue | 🟠 MajorOnly count upstream/server failures against the circuit breaker.
This currently trips the breaker for every non-2xx, so repeated client 4xxs can evict a healthy Gemini worker. The JSON parse branch is also an upstream contract failure and should surface as 502, not
internal_error.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/gemini/steps/non_stream_execution.rs` around lines 87 - 105, Only record circuit breaker failures for upstream/server errors and treat JSON parse failures as upstream contract failures: change the response.status() branch so worker.circuit_breaker().record_failure() is called only when status.is_server_error() (5xx) while still returning non-2xx statuses to the client; in the response.json().await Err(e) branch treat it as a 502 upstream error (use StatusCode::BAD_GATEWAY) and call worker.circuit_breaker().record_failure() before returning an error response instead of using error::internal_error("parse_error", ...), so parsing failures count against the breaker and client 4xx responses do not.
63-68:⚠️ Potential issue | 🟠 MajorDon’t let URL sniffing choose Gemini auth headers.
apply_provider_headers()derives the provider fromupstream_url, so a Gemini worker behindlocalhost, a proxy, or a custom domain is treated as Generic and getsAuthorizationinstead ofx-goog-api-key. The mock server explicitly tolerates that fallback, so tests still pass, but real proxied Gemini backends will fail auth.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/gemini/steps/non_stream_execution.rs` around lines 63 - 68, apply_provider_headers currently sniffs provider from upstream_url which causes Gemini workers behind proxies/domains to be treated as Generic and send Authorization instead of x-goog-api-key; change the call site in non_stream_execution.rs so provider detection is explicit: after computing auth_header via extract_gemini_auth_header(ctx.input.headers.as_ref(), worker.api_key()), call the variant of apply_provider_headers that accepts an explicit provider hint (or add one) and pass the Gemini/provider-for-api-key enum so the request built by ctx.components.client.post(upstream_url).json(&payload) will use the x-goog-api-key header; if no such API exists add a short overload to apply_provider_headers to accept a Provider enum (or a flag) and ensure it prefers the API-key header when auth_header is an API key.model_gateway/tests/common/mock_gemini_server.rs (1)
181-207:⚠️ Potential issue | 🟠 MajorMake the mock
/v1beta/modelsendpoint match Gemini’s real contract.Production discovery expects Gemini auth plus
{ "models": [{ "name": "models/..." }] }, but this mock still returns OpenAI-styledata[].idand skipscheck_auth(). Any test that exercises refresh/discovery is validating the wrong protocol.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/tests/common/mock_gemini_server.rs` around lines 181 - 207, The mock_models handler is returning OpenAI-style data[].id and not enforcing Gemini auth; update the mock /v1beta/models implementation (mock_models) to call the existing check_auth() at the start and return the Gemini-shaped JSON: top-level "models": [ { "name": "models/<model-id>" }, ... ] (e.g., "models/gemini-2.5-flash", "models/gemini-2.5-pro", "models/deep-research-pro-preview-12-2025") instead of data[].id so tests validate the real Gemini discovery contract.model_gateway/src/core/steps/worker/external/discover_models.rs (1)
232-245:⚠️ Potential issue | 🟠 MajorDon’t decide Gemini discovery from the URL alone.
execute()still derivesproviderfromProviderType::from_url(&config.url), so Gemini backends behindlocalhost, a proxy, or any custom domain fall through to/v1/modelsplus Bearer auth. That makes the new/v1beta/modelsandx-goog-api-keypath unusable for proxied Gemini deployments.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/steps/worker/external/discover_models.rs` around lines 232 - 245, The code chooses Gemini based solely on ProviderType::from_url(&config.url) inside execute(), causing proxied/custom domains to miss the /v1beta/models path and x-goog-api-key header; change the detection to prefer an explicit provider indicator from the config (e.g., a config.provider or config.api_type field) when available and only fall back to ProviderType::from_url(&config.url), then use that resolved provider variable for models_path and header selection (the logic that sets models_path/models_url and the header branches that call uses_x_api_key() or check for ProviderType::Gemini) so proxied Gemini deployments are handled correctly.model_gateway/src/routers/header_utils.rs (1)
177-183:⚠️ Potential issue | 🟡 MinorNormalize the Bearer scheme case-insensitively before setting
x-goog-api-key.
strip_prefix("Bearer ")leaves lowercase or mixed-case schemes in the forwarded key, soauthorization: bearer ...reaches Gemini as an invalid API key.Suggested fix
- let api_key = auth_str.strip_prefix("Bearer ").unwrap_or(auth_str); - req = req.header("x-goog-api-key", api_key); + let api_key = auth_str + .split_once(' ') + .filter(|(scheme, token)| { + scheme.eq_ignore_ascii_case("bearer") && !token.trim().is_empty() + }) + .map(|(_, token)| token.trim()) + .unwrap_or(auth_str.trim()); + if !api_key.is_empty() { + req = req.header("x-goog-api-key", api_key); + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/header_utils.rs` around lines 177 - 183, In the ApiProvider::Gemini branch where you read auth_header and compute api_key, handle the "Bearer " scheme case-insensitively: parse auth.to_str() into tokens (e.g., splitn on whitespace) and if the first token equals "bearer" ignoring case, use the second token as api_key; otherwise use the entire auth_str; then set req.header("x-goog-api-key", api_key). This ensures mixed- or lower-case "bearer" schemes are normalized before forwarding to Gemini.model_gateway/src/main.rs (1)
1033-1035: 🛠️ Refactor suggestion | 🟠 MajorKeep all HTTP-only external backends explicit in the
connection_modematch.This branch now special-cases OpenAI and Gemini but still leaves Anthropic to the fallback URL heuristic. It works today, but the external-backend invariant is being encoded here, so Anthropic should be included too.
Suggested fix
- RoutingMode::OpenAI { .. } | RoutingMode::Gemini { .. } => ConnectionMode::Http, + RoutingMode::OpenAI { .. } + | RoutingMode::Anthropic { .. } + | RoutingMode::Gemini { .. } => ConnectionMode::Http,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/main.rs` around lines 1033 - 1035, The match for computing connection_mode in main.rs currently special-cases RoutingMode::OpenAI and RoutingMode::Gemini to ConnectionMode::Http but leaves Anthropic to the fallback heuristic; update the match expression inside wherever connection_mode is computed (the match on &mode that returns ConnectionMode::Http or Self::determine_connection_mode(&all_urls)) to also explicitly match RoutingMode::Anthropic { .. } => ConnectionMode::Http so Anthropic is treated as an HTTP-only external backend alongside OpenAI and Gemini.model_gateway/src/core/job_queue.rs (1)
706-718:⚠️ Potential issue | 🟠 MajorCarry the Gemini provider in the worker config, not just in log strings.
submit_external_worker_jobsstill builds the same generic externalWorkerSpecfor Gemini as for OpenAI/Anthropic. That means any later provider-specific behavior has to re-infer Gemini from the URL, and this repo's current heuristic only recognizesgoogleapis.com; Gemini proxies, custom hosts, and localhost workers will miss Gemini discovery/auth handling.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/job_queue.rs` around lines 706 - 718, submit_external_worker_jobs currently constructs a generic external WorkerSpec via build_external_worker_config without preserving the provider_name, so Gemini must not be inferred from URL strings; update the code to pass provider_name into build_external_worker_config (or extend the config/WorkerSpec constructor) and ensure the returned WorkerSpec includes a provider field set to provider_name (e.g., "gemini", "openai", "anthropic"); also update any downstream logic that checks the URL to instead read the WorkerSpec.provider to drive provider-specific auth/behavior in functions that consume WorkerSpec.model_gateway/src/routers/gemini/utils.rs (1)
10-18:⚠️ Potential issue | 🟠 MajorHonor Gemini auth precedence here: non-empty
x-goog-api-key→Authorization: Bearer→ worker key.This helper only reads
x-goog-api-key, so the same BYOK request can be accepted inmodel_gateway/src/routers/header_utils.rsbut treated as unauthenticated in Gemini worker-selection/request-building paths. It also lets an emptyx-goog-api-keysuppress the worker-key fallback.Suggested fix
pub(crate) fn extract_gemini_auth_header( headers: Option<&HeaderMap>, worker_api_key: Option<&String>, ) -> Option<HeaderValue> { - // Passthrough: try user's x-goog-api-key header first - let user_auth = headers.and_then(|h| h.get("x-goog-api-key").cloned()); - - // Return user's key if provided, otherwise use worker's API key as-is (no Bearer prefix) - user_auth.or_else(|| worker_api_key.and_then(|k| HeaderValue::from_str(k).ok())) + let user_auth = headers.and_then(|h| { + h.get("x-goog-api-key") + .and_then(|v| v.to_str().ok().map(str::trim)) + .filter(|v| !v.is_empty()) + .map(|v| v.to_owned()) + .or_else(|| { + h.get("authorization") + .and_then(|v| v.to_str().ok()) + .and_then(|auth| auth.split_once(' ')) + .filter(|(scheme, token)| { + scheme.eq_ignore_ascii_case("bearer") && !token.trim().is_empty() + }) + .map(|(_, token)| token.trim().to_owned()) + }) + }); + + user_auth + .and_then(|key| HeaderValue::from_str(&key).ok()) + .or_else(|| worker_api_key.and_then(|k| HeaderValue::from_str(k).ok())) }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/gemini/utils.rs` around lines 10 - 18, extract_gemini_auth_header currently only reads x-goog-api-key and treats an empty value as present, which suppresses falling back to Authorization: Bearer or the worker key; update the function to (1) prefer a non-empty "x-goog-api-key" header, (2) if that is absent or empty, check the "authorization" header and use it if it contains a Bearer token, and (3) otherwise fall back to worker_api_key (as a HeaderValue without adding "Bearer"). Ensure you test for empty header values when reading headers.and_then(...).cloned() and parse the "authorization" header string to accept only "Bearer <token>" before returning it.model_gateway/tests/api/interactions_api_test.rs (2)
165-170:⚠️ Potential issue | 🟠 MajorKeep the missing-model assertion strict.
With a healthy mock upstream, this path should be
404 Not Found. Allowing503 Service Unavailablemasks regressions in the miss/refresh flow and lets the wrong failure mode pass.🧪 Tighten the assertion
- assert!( - response.status() == StatusCode::NOT_FOUND - || response.status() == StatusCode::SERVICE_UNAVAILABLE, - "Expected 404 or 503, got {}", - response.status() - ); + assert_eq!(response.status(), StatusCode::NOT_FOUND);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/tests/api/interactions_api_test.rs` around lines 165 - 170, The test currently allows either 404 or 503, which hides regressions; update the assertion in the interactions test to require only StatusCode::NOT_FOUND by removing the OR branch and asserting response.status() == StatusCode::NOT_FOUND (and keep the same failure message), so the missing-model path fails if a 503 is returned; locate the assertion using response.status() and StatusCode::NOT_FOUND in interactions_api_test.rs and tighten it accordingly.
277-306:⚠️ Potential issue | 🟠 MajorThis still doesn't prove Gemini auth was forwarded as
x-goog-api-key.The test itself notes that localhost takes the generic header path and that the mock accepts both
x-goog-api-keyandAuthorization. A regression back to bearer-only forwarding would still pass. Make the mock capture or reject the header shape so this assertion proves the Gemini contract.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/tests/api/interactions_api_test.rs` around lines 277 - 306, The test currently can pass even if only Authorization was forwarded; modify the test and mock so it verifies the exact header shape: update MockGeminiServer (constructed via MockGeminiServer::new_with_auth) to record incoming request headers and fail or return 400 if Authorization bearer is present instead of x-goog-api-key, then in test_interactions_forwards_api_key_header assert the mock recorded "x-goog-api-key" == "test-key-123" (and assert that "authorization" is absent) after calling router.route_interactions; alternatively add an explicit mock expectation that only x-goog-api-key is accepted rather than accepting both.model_gateway/src/routers/gemini/steps/request_building.rs (1)
88-97:⚠️ Potential issue | 🟠 Major
model_idis overloaded, so agent requests with a path override can still diverge.At this point
ctx.input.model_idcan mean either “explicit URL override” or just “the resolved body target”. When anagentrequest comes through a model-path route, worker selection already routed onmodel_id, but this transform leaves the payload targetingagent, so the chosen worker and forwarded target can disagree. Split those into separate fields, or rejectagenton model-path routes before this step.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/gemini/steps/request_building.rs` around lines 88 - 97, The current logic uses ctx.input.model_id for both “explicit model-path URL override” and the resolved body target, causing agent requests routed via a model-path to be left targeting agent while the worker was selected by model_id; update the code to distinguish these cases: add or use a separate field (e.g. model_path_override or resolved_model_id) so that when populating the outgoing body (the obj.insert("model", ...) branch) you only inject "model" when the model was explicitly intended for the body (resolved_model_id) and not when model_path_override was used for routing, or alternatively reject requests where original_request.agent.is_some() while a model-path override is present before this transform; ensure checks reference ctx.input.model_id (or the new model_path_override/resolved_model_id), ctx.input.original_request.model, and ctx.input.original_request.agent to determine the correct behavior.
🤖 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/steps/worker_selection.rs`:
- Around line 49-54: The code currently calls
extract_gemini_auth_header(ctx.input.headers.as_ref(), None) which forces no
worker-local API key and causes select_worker (which may invoke
refresh_external_models) to fetch /v1beta/models unauthenticated; change the
call so the second argument supplies the worker-specific API key when available
(derive the per-worker x-goog-api-key from the worker configuration or selection
context and pass it into extract_gemini_auth_header) so select_worker and any
refresh_external_models calls use per-worker auth rather than always None.
In `@model_gateway/src/server.rs`:
- Around line 247-250: The current code sets model_id =
body.model.as_deref().or(body.agent.as_deref()), which incorrectly treats
body.agent as a model id; change this to only use the explicit model field
(e.g., model_id = body.model.as_deref()) and return a 400 Bad Request when model
is missing instead of falling back to body.agent, leaving any future agent→model
resolution to an explicit separate step; update the call to
state.router.route_interactions(...) to receive the validated model_id (or
short-circuit with the 400) so agent identifiers are never passed into
route_interactions as a model id.
---
Duplicate comments:
In `@model_gateway/src/core/job_queue.rs`:
- Around line 706-718: submit_external_worker_jobs currently constructs a
generic external WorkerSpec via build_external_worker_config without preserving
the provider_name, so Gemini must not be inferred from URL strings; update the
code to pass provider_name into build_external_worker_config (or extend the
config/WorkerSpec constructor) and ensure the returned WorkerSpec includes a
provider field set to provider_name (e.g., "gemini", "openai", "anthropic");
also update any downstream logic that checks the URL to instead read the
WorkerSpec.provider to drive provider-specific auth/behavior in functions that
consume WorkerSpec.
In `@model_gateway/src/core/steps/worker/external/discover_models.rs`:
- Around line 232-245: The code chooses Gemini based solely on
ProviderType::from_url(&config.url) inside execute(), causing proxied/custom
domains to miss the /v1beta/models path and x-goog-api-key header; change the
detection to prefer an explicit provider indicator from the config (e.g., a
config.provider or config.api_type field) when available and only fall back to
ProviderType::from_url(&config.url), then use that resolved provider variable
for models_path and header selection (the logic that sets models_path/models_url
and the header branches that call uses_x_api_key() or check for
ProviderType::Gemini) so proxied Gemini deployments are handled correctly.
In `@model_gateway/src/main.rs`:
- Around line 1033-1035: The match for computing connection_mode in main.rs
currently special-cases RoutingMode::OpenAI and RoutingMode::Gemini to
ConnectionMode::Http but leaves Anthropic to the fallback heuristic; update the
match expression inside wherever connection_mode is computed (the match on &mode
that returns ConnectionMode::Http or Self::determine_connection_mode(&all_urls))
to also explicitly match RoutingMode::Anthropic { .. } => ConnectionMode::Http
so Anthropic is treated as an HTTP-only external backend alongside OpenAI and
Gemini.
In `@model_gateway/src/routers/gemini/steps/non_stream_execution.rs`:
- Around line 87-105: Only record circuit breaker failures for upstream/server
errors and treat JSON parse failures as upstream contract failures: change the
response.status() branch so worker.circuit_breaker().record_failure() is called
only when status.is_server_error() (5xx) while still returning non-2xx statuses
to the client; in the response.json().await Err(e) branch treat it as a 502
upstream error (use StatusCode::BAD_GATEWAY) and call
worker.circuit_breaker().record_failure() before returning an error response
instead of using error::internal_error("parse_error", ...), so parsing failures
count against the breaker and client 4xx responses do not.
- Around line 63-68: apply_provider_headers currently sniffs provider from
upstream_url which causes Gemini workers behind proxies/domains to be treated as
Generic and send Authorization instead of x-goog-api-key; change the call site
in non_stream_execution.rs so provider detection is explicit: after computing
auth_header via extract_gemini_auth_header(ctx.input.headers.as_ref(),
worker.api_key()), call the variant of apply_provider_headers that accepts an
explicit provider hint (or add one) and pass the Gemini/provider-for-api-key
enum so the request built by
ctx.components.client.post(upstream_url).json(&payload) will use the
x-goog-api-key header; if no such API exists add a short overload to
apply_provider_headers to accept a Provider enum (or a flag) and ensure it
prefers the API-key header when auth_header is an API key.
In `@model_gateway/src/routers/gemini/steps/request_building.rs`:
- Around line 88-97: The current logic uses ctx.input.model_id for both
“explicit model-path URL override” and the resolved body target, causing agent
requests routed via a model-path to be left targeting agent while the worker was
selected by model_id; update the code to distinguish these cases: add or use a
separate field (e.g. model_path_override or resolved_model_id) so that when
populating the outgoing body (the obj.insert("model", ...) branch) you only
inject "model" when the model was explicitly intended for the body
(resolved_model_id) and not when model_path_override was used for routing, or
alternatively reject requests where original_request.agent.is_some() while a
model-path override is present before this transform; ensure checks reference
ctx.input.model_id (or the new model_path_override/resolved_model_id),
ctx.input.original_request.model, and ctx.input.original_request.agent to
determine the correct behavior.
In `@model_gateway/src/routers/gemini/utils.rs`:
- Around line 10-18: extract_gemini_auth_header currently only reads
x-goog-api-key and treats an empty value as present, which suppresses falling
back to Authorization: Bearer or the worker key; update the function to (1)
prefer a non-empty "x-goog-api-key" header, (2) if that is absent or empty,
check the "authorization" header and use it if it contains a Bearer token, and
(3) otherwise fall back to worker_api_key (as a HeaderValue without adding
"Bearer"). Ensure you test for empty header values when reading
headers.and_then(...).cloned() and parse the "authorization" header string to
accept only "Bearer <token>" before returning it.
In `@model_gateway/src/routers/header_utils.rs`:
- Around line 177-183: In the ApiProvider::Gemini branch where you read
auth_header and compute api_key, handle the "Bearer " scheme case-insensitively:
parse auth.to_str() into tokens (e.g., splitn on whitespace) and if the first
token equals "bearer" ignoring case, use the second token as api_key; otherwise
use the entire auth_str; then set req.header("x-goog-api-key", api_key). This
ensures mixed- or lower-case "bearer" schemes are normalized before forwarding
to Gemini.
In `@model_gateway/tests/api/interactions_api_test.rs`:
- Around line 165-170: The test currently allows either 404 or 503, which hides
regressions; update the assertion in the interactions test to require only
StatusCode::NOT_FOUND by removing the OR branch and asserting response.status()
== StatusCode::NOT_FOUND (and keep the same failure message), so the
missing-model path fails if a 503 is returned; locate the assertion using
response.status() and StatusCode::NOT_FOUND in interactions_api_test.rs and
tighten it accordingly.
- Around line 277-306: The test currently can pass even if only Authorization
was forwarded; modify the test and mock so it verifies the exact header shape:
update MockGeminiServer (constructed via MockGeminiServer::new_with_auth) to
record incoming request headers and fail or return 400 if Authorization bearer
is present instead of x-goog-api-key, then in
test_interactions_forwards_api_key_header assert the mock recorded
"x-goog-api-key" == "test-key-123" (and assert that "authorization" is absent)
after calling router.route_interactions; alternatively add an explicit mock
expectation that only x-goog-api-key is accepted rather than accepting both.
In `@model_gateway/tests/common/mock_gemini_server.rs`:
- Around line 181-207: The mock_models handler is returning OpenAI-style
data[].id and not enforcing Gemini auth; update the mock /v1beta/models
implementation (mock_models) to call the existing check_auth() at the start and
return the Gemini-shaped JSON: top-level "models": [ { "name":
"models/<model-id>" }, ... ] (e.g., "models/gemini-2.5-flash",
"models/gemini-2.5-pro", "models/deep-research-pro-preview-12-2025") instead of
data[].id so tests validate the real Gemini discovery contract.
In `@model_gateway/tests/common/mod.rs`:
- Around line 412-429: create_test_context() seeds Gemini external workers but
create_test_context_with_parsers() and create_test_context_with_mcp_config()
still only register OpenAI workers; update both helpers to mirror the Gemini
worker registration logic: detect RoutingMode::Gemini and for each url build the
same Vec<ModelCard> and Arc<dyn Worker> via
BasicWorkerBuilder::new(url).worker_type(WorkerType::Regular).runtime_type(RuntimeType::External).models(models).build()
and call app_context.worker_registry.register(worker) so the registry is
populated the same way as in create_test_context().
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: b7438bf5-df71-4e4c-b606-1d7c4086f449
📒 Files selected for processing (25)
crates/protocols/src/interactions.rsmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/config/validation.rsmodel_gateway/src/core/job_queue.rsmodel_gateway/src/core/steps/worker/external/discover_models.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/factory.rsmodel_gateway/src/routers/gemini/context.rsmodel_gateway/src/routers/gemini/mod.rsmodel_gateway/src/routers/gemini/router.rsmodel_gateway/src/routers/gemini/state.rsmodel_gateway/src/routers/gemini/steps/non_stream_execution.rsmodel_gateway/src/routers/gemini/steps/request_building.rsmodel_gateway/src/routers/gemini/steps/response_processing.rsmodel_gateway/src/routers/gemini/steps/worker_selection.rsmodel_gateway/src/routers/gemini/utils.rsmodel_gateway/src/routers/header_utils.rsmodel_gateway/src/routers/mod.rsmodel_gateway/src/routers/router_manager.rsmodel_gateway/src/server.rsmodel_gateway/tests/api/interactions_api_test.rsmodel_gateway/tests/api/mod.rsmodel_gateway/tests/common/mock_gemini_server.rsmodel_gateway/tests/common/mod.rs
💤 Files with no reviewable changes (1)
- model_gateway/src/routers/gemini/state.rs
2 of 3 tasks
Contributor
|
Hi @XinyueZhang369, 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
Signed-off-by: XinyueZhang369 <zoeyzhang369@gmail.com>
XinyueZhang369
force-pushed
the
xz/non-stream-interactions-request
branch
from
March 9, 2026 23:22
7784367 to
7b0d57b
Compare
Collaborator
Author
Thanks for reviewing! |
mergify
Bot
removed
the
needs-rebase
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7b0d57b569
ℹ️ 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 (3)
model_gateway/src/server.rs (1)
242-252:⚠️ Potential issue | 🟠 MajorDon't treat
agentas a model id in the phase-1 interactions path.
body.model.as_deref().or(body.agent.as_deref())lets agent-only requests route on an agent identifier instead of failing fast for a missing model. That both hides the expected 400 path and can send an invalid id into worker selection.Suggested fix
- let model_id = body.model.as_deref().or(body.agent.as_deref()); + let model_id = body.model.as_deref();🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/server.rs` around lines 242 - 252, In v1_interactions, don't use body.agent as a fallback for model id; change model_id to only use body.model.as_deref() (remove .or(body.agent.as_deref())), and if model is None return a 400 error immediately (instead of calling state.router.route_interactions) so agent-only requests fail fast; update the logic in v1_interactions to perform this check before calling state.router.route_interactions.model_gateway/src/routers/router_manager.rs (1)
614-626:⚠️ Potential issue | 🟠 MajorResolve the model before selecting a router in IGW mode.
route_interactionsis the only IGW request path here that skipsresolve_model_id. When the request omits all model selectors, this falls through toselect_router_for_request(headers, None), so single-model deployments won’t infer the only model and multi-model deployments won’t return the expected 400; the request can land on an arbitrary/default router instead.Proposed fix
async fn route_interactions( &self, headers: Option<&HeaderMap>, body: &InteractionsRequest, model_id: Option<&str>, ) -> Response { let selected_model = model_id.or(body.model.as_deref()).or(body.agent.as_deref()); - let router = self.select_router_for_request(headers, selected_model); + let effective_model_id = if self.enable_igw { + match self.resolve_model_id(selected_model) { + Ok(id) => Some(id), + Err(err_response) => return *err_response, + } + } else { + None + }; + + let effective_model = effective_model_id.as_deref().or(selected_model); + let router = self.select_router_for_request(headers, effective_model); if let Some(router) = router { router - .route_interactions(headers, body, selected_model) + .route_interactions(headers, body, effective_model) .await } else { ( StatusCode::NOT_FOUND, "No router available to handle interactions request",🤖 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 614 - 626, route_interactions currently calls select_router_for_request with selected_model possibly None, skipping resolve_model_id; change it to first call resolve_model_id(headers, body.model.as_deref(), body.agent.as_deref()) (or the existing resolve_model_id helper) to obtain the effective model id (and handle/return any error Response it may produce), then pass that resolved model id into select_router_for_request and into router.route_interactions; update references to selected_model so they use the resolved id rather than the original optional chain.model_gateway/src/core/job_queue.rs (1)
706-718:⚠️ Potential issue | 🟠 MajorCarry the provider through the AddWorker config, not just the log strings.
provider_nameis only used for messages here. The helper still builds a provider-agnosticWorkerSpec, so if the registration/discovery path still infers provider fromconfig.url, Gemini endpoints behind proxies/custom hostnames will regress to the generic/OpenAI path. That also breaks the new provider-aware branch inmodel_gateway/src/routers/router_manager.rs(Lines 232-236) whenprovider_for_model()is missing.Please verify that an explicit provider/routing-mode field is threaded from
build_external_worker_configthrough workflow data into discovery, and that discovery prefers it over URL heuristics. If the only Gemini signal is still URL-based (/v1beta/models,x-goog-api-key, hostname checks), this issue is still present.#!/bin/bash set -euo pipefail echo "== WorkerSpec fields ==" fd 'worker.rs$' . -x rg -n -C3 --type rust 'pub struct WorkerSpec|provider|routing_mode' {} echo echo "== External worker config and workflow data ==" rg -n -C4 --type rust 'build_external_worker_config\s*\(|submit_external_worker_jobs\s*\(|create_worker_workflow_data\s*\(' . echo echo "== Discovery/provider selection ==" fd 'discover_models.rs$' . -x rg -n -C4 --type rust 'provider|routing_mode|v1beta/models|x-goog-api-key|url' {}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/core/job_queue.rs` around lines 706 - 718, submit_external_worker_jobs currently only passes provider_name into logs while build_external_worker_config produces a provider-agnostic config, so thread an explicit provider/routing_mode through the config and into the WorkerSpec and workflow data creation: update build_external_worker_config to accept/return a provider or routing_mode field, ensure create_worker_workflow_data copies that provider/routing_mode into the workflow payload, and make the discovery/provider selection (used by provider_for_model/router_manager) prefer this explicit field over URL heuristics (v1beta paths, x-goog-api-key, hostname checks); touch submit_external_worker_jobs, build_external_worker_config, create_worker_workflow_data, WorkerSpec, and the discovery/provider_for_model logic to propagate and prefer the explicit 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/server.rs`:
- Line 597: The server currently only mounts "/v1/interactions" so path-based
model overrides never reach route_interactions; add a second route for
"/v1/interactions/{model_id}" and implement a handler (e.g.,
v1_interactions_with_model) that extracts the path parameter (model_id) and
forwards it into the same logic used by route_interactions (or call
route_interactions with the extracted model_id) so the route_interactions code
can apply the model-id override; update the route registration to include
post(v1_interactions_with_model) alongside the existing post(v1_interactions).
---
Duplicate comments:
In `@model_gateway/src/core/job_queue.rs`:
- Around line 706-718: submit_external_worker_jobs currently only passes
provider_name into logs while build_external_worker_config produces a
provider-agnostic config, so thread an explicit provider/routing_mode through
the config and into the WorkerSpec and workflow data creation: update
build_external_worker_config to accept/return a provider or routing_mode field,
ensure create_worker_workflow_data copies that provider/routing_mode into the
workflow payload, and make the discovery/provider selection (used by
provider_for_model/router_manager) prefer this explicit field over URL
heuristics (v1beta paths, x-goog-api-key, hostname checks); touch
submit_external_worker_jobs, build_external_worker_config,
create_worker_workflow_data, WorkerSpec, and the discovery/provider_for_model
logic to propagate and prefer the explicit provider.
In `@model_gateway/src/routers/router_manager.rs`:
- Around line 614-626: route_interactions currently calls
select_router_for_request with selected_model possibly None, skipping
resolve_model_id; change it to first call resolve_model_id(headers,
body.model.as_deref(), body.agent.as_deref()) (or the existing resolve_model_id
helper) to obtain the effective model id (and handle/return any error Response
it may produce), then pass that resolved model id into select_router_for_request
and into router.route_interactions; update references to selected_model so they
use the resolved id rather than the original optional chain.
In `@model_gateway/src/server.rs`:
- Around line 242-252: In v1_interactions, don't use body.agent as a fallback
for model id; change model_id to only use body.model.as_deref() (remove
.or(body.agent.as_deref())), and if model is None return a 400 error immediately
(instead of calling state.router.route_interactions) so agent-only requests fail
fast; update the logic in v1_interactions to perform this check before calling
state.router.route_interactions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: cbc7362e-1aaa-40af-b2ec-9feb6d6aee61
📒 Files selected for processing (12)
crates/protocols/src/interactions.rsmodel_gateway/src/config/builder.rsmodel_gateway/src/config/types.rsmodel_gateway/src/config/validation.rsmodel_gateway/src/core/job_queue.rsmodel_gateway/src/main.rsmodel_gateway/src/routers/factory.rsmodel_gateway/src/routers/gemini/context.rsmodel_gateway/src/routers/gemini/router.rsmodel_gateway/src/routers/mod.rsmodel_gateway/src/routers/router_manager.rsmodel_gateway/src/server.rs
XinyueZhang369
changed the title
slin1237
approved these changes
Mar 10, 2026
CatherineSue
approved these changes
Mar 10, 2026
6 tasks
Merged
3 tasks
2 of 3 tasks
4 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Problem
The model gateway currently supports OpenAI and Anthropic as external provider backends, but lacks support for Google's Gemini API. Users who want to route requests through Gemini-compatible endpoints cannot do so.
Solution
Add the routing infrastructure and registration plumbing for Gemini as a new routing mode in the model gateway. This PR wires up the config, CLI, job queue, router factory, server endpoint, and router manager so that the
Gemini router is fully registered and routable. The router implementation stubs (worker selection, request building, upstream execution, response processing) and integration tests will follow in a subsequent PR. Refer to #417 for more details of the gemini router structure
Changes
RoutingMode::Geminivariant to config types, validation, CLI--backend geminiflag, and job queue worker initializationsubmit_external_worker_jobshelper, add Gemini to the combined match armGeminiRouterin the router factory (both single-router and IGW multi-router modes), addHTTP_GEMINIrouter ID, include increate_all_routersPOST /v1/interactionsroute withValidatedJson<InteractionsRequest>extractorroute_interactionsmethod (returns 501) so existing routers are unaffectedroute_interactionsdispatch and provider-aware routing inget_router_for_modelusingProviderTypeAppContextwithSharedComponents(client, worker_registry, mcp_orchestrator, request_timeout), implementRouterTrait::route_interactionsNormalizableforInteractionsRequestforValidatedJsonextractor compatibilityTest Plan
cargo buildsucceeds with new Gemini config variantscargo run -p smg --bin smg -- --backend gemini --worker-urls "https://generativelanguage.googleapis.com" --port 8080to start smg for gemini backend. The router starts successfully and return 501 not implemented for /v1/interactions requestChecklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
New Features
Tests