refactor(gateway): split OpenAI router.rs into chat and health modules - #726
Conversation
Extract route_chat and health/server-info endpoints from the monolithic router.rs (~988 lines) into focused modules, reducing it to a thin dispatcher (~680 lines). What changed: - model_gateway/src/routers/openai/chat.rs: new module containing ChatDeps struct and route_chat free function, with the full chat completion routing logic (worker selection, provider transform, retry loop with streaming/non-streaming handling) - model_gateway/src/routers/openai/health.rs: new module containing health_generate and get_server_info free functions, with the external_workers helper inlined - model_gateway/src/routers/openai/mod.rs: added mod chat and mod health - model_gateway/src/routers/openai/router.rs: removed route_chat body (~250 lines) and health/server-info bodies (~60 lines), RouterTrait methods now delegate to the extracted modules. Extracted resolve_provider as pub(super) free function shared by chat.rs and route_responses. Removed unused shared_components() accessor. Why: router.rs mixed chat dispatch, health endpoints, responses orchestration, storage queries, and realtime methods. Splitting improves navigability and sets up further extractions (responses in PR 4, realtime in PR 5). How: Uses free functions with dependency structs (ChatDeps) rather than splitting impl blocks across files, since OpenAIRouter fields are private. This matches the pattern used in Anthropic/Gemini routers (e.g. streaming::execute(&router_ctx, req_ctx)). Pure refactor — no behavior changes. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
github-actions
Bot
added
model-gateway
📝 WalkthroughWalkthroughAdds a new OpenAI chat routing module that forwards chat completion requests to upstream workers with retries, provider-specific transformations, streaming and non-streaming response handling, extensive error metrics, and a health module that reports worker/registry status. Router refactoring delegates provider resolution and chat/health flows to new modules. Changes
Sequence DiagramsequenceDiagram
participant Client as Client
participant Router as Router (route_chat)
participant Selector as WorkerSelector
participant Provider as Provider
participant Worker as Upstream Worker
participant Retry as RetryExecutor
Client->>Router: ChatCompletionRequest
Router->>Router: record incoming metrics
Router->>Selector: select worker
Selector-->>Router: selected worker
Router->>Provider: resolve & apply transformations
Provider-->>Router: transformed payload
Router->>Retry: prepare RequestContext & start retries
loop retry attempts
Retry->>Worker: POST /chat/completions (with headers)
Worker-->>Retry: Response (stream or body)
end
alt streaming
Retry->>Client: relay event stream (preserve content-type/status)
else non-streaming
Retry->>Client: forward body (preserve content-type/status)
end
Router->>Router: update circuit breaker & record final metrics
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
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 refactors the OpenAI router by splitting its monolithic 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
|
![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 refactors router.rs into chat.rs and health.rs, improving maintainability and readability. However, it highlights existing security concerns: unvalidated user input for model names in metrics can lead to Denial of Service, and internal worker URLs are exposed in health and server info endpoints, causing information leakage. Additionally, there are opportunities to improve efficiency in the new health.rs module by reducing redundant iterations and allocations.
![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/openai/chat.rs`:
- Around line 48-50: The code resolves an override via model_id into the local
variable model but never rewrites the outgoing chat request or the
RequestContext, causing mismatched metrics/provider selection vs the upstream
payload; update the chat path to mutate or build a new request_body where
body.model is set to model (i.e., apply the model_id override into the
serialized payload used for routing/worker lookup) before any provider
resolution and before serializing the body, and pass request_body.clone() into
RequestContext::for_chat so downstream state and metrics use the same overridden
model; ensure all branches that inspect body.model or serialize body (the
selection logic around model usage and the serialization near where
streaming/worker lookup occurs) use this rewritten request_body.
In `@model_gateway/src/routers/openai/health.rs`:
- Around line 60-68: The response currently returns JSON as plain text; change
the return to use axum's JSON wrapper so the Content-Type is application/json.
Replace the final line that returns (StatusCode::OK,
info.to_string()).into_response() with (StatusCode::OK,
axum::Json(info)).into_response() (and add use axum::Json if not already
imported); ensure the `info` value remains a serde_json::Value (from json!) so
axum::Json can serialize it.
In `@model_gateway/src/routers/openai/router.rs`:
- Around line 56-70: The helper resolve_provider currently only checks the
optional override (model_id) and skips provider resolution when model_id is
None, causing non-default models to fall back to default_provider_arc; update
the call sites (this file where resolve_provider is invoked and
model_gateway/src/routers/openai/chat.rs around the noted lines) to pass the
already-computed effective model string instead of the raw optional override,
and adjust resolve_provider (and its signature if needed) to use that effective
model (so its logic still uses worker.provider_for_model(...) and
ProviderType::from_model_name(...) on the effective model) rather than relying
on an Option that may be None. Ensure resolve_provider no longer ignores the
effective request model and that both callers supply that resolved model value.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: ce473558-e432-42f5-ab8a-bf3e08abb9eb
📒 Files selected for processing (4)
model_gateway/src/routers/openai/chat.rsmodel_gateway/src/routers/openai/health.rsmodel_gateway/src/routers/openai/mod.rsmodel_gateway/src/routers/openai/router.rs
| let start = Instant::now(); | ||
| let model = model_id.unwrap_or(body.model.as_str()); | ||
| let streaming = body.stream; |
There was a problem hiding this comment.
Rewrite the chat request with the override model before routing it.
The function resolves model from model_id, but Line 64 still selects on body.model and Line 85 serializes the original body unchanged. If model_id is set, metrics/provider resolution point at one model while worker lookup and the upstream payload use another.
🐛 Proposed fix
- let start = Instant::now();
- let model = model_id.unwrap_or(body.model.as_str());
+ let start = Instant::now();
+ let mut request_body = body.clone();
+ if let Some(model_id) = model_id {
+ request_body.model = model_id.to_string();
+ }
+ let model = request_body.model.as_str();
let streaming = body.stream;
@@
.select_worker(&SelectWorkerRequest {
- model_id: body.model.as_str(),
+ model_id: request_body.model.as_str(),
headers,
provider: Some(ProviderType::OpenAI),
..Default::default()
})
@@
- let mut payload = match to_value(body) {
+ let mut payload = match to_value(&request_body) {
Ok(v) => v,I'd also pass request_body.clone() into RequestContext::for_chat so downstream state stays aligned.
Also applies to: 61-68, 85-120
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/routers/openai/chat.rs` around lines 48 - 50, The code
resolves an override via model_id into the local variable model but never
rewrites the outgoing chat request or the RequestContext, causing mismatched
metrics/provider selection vs the upstream payload; update the chat path to
mutate or build a new request_body where body.model is set to model (i.e., apply
the model_id override into the serialized payload used for routing/worker
lookup) before any provider resolution and before serializing the body, and pass
request_body.clone() into RequestContext::for_chat so downstream state and
metrics use the same overridden model; ensure all branches that inspect
body.model or serialize body (the selection logic around model usage and the
serialization near where streaming/worker lookup occurs) use this rewritten
request_body.
What changed: - model_gateway/src/routers/openai/router.rs: resolve_provider now takes &str instead of Option<&str>, so provider resolution always uses the effective model (not just the URL override). route_responses call site updated to pass the computed model. - model_gateway/src/routers/openai/chat.rs: worker selection now uses the effective model (model_id override or body.model) instead of always using body.model. When model_id overrides the body model, the serialized payload is patched to match. resolve_provider call updated to pass &str. - model_gateway/src/routers/openai/health.rs: get_server_info now returns application/json Content-Type via axum::Json instead of text/plain. Why: - resolve_provider with Option<&str> skipped provider lookup entirely for body-driven requests (model_id=None), always returning default_provider. Passing the effective model ensures correct provider selection. - Worker selection using body.model while metrics/provider used model_id caused inconsistent routing when a URL path override was provided. - get_server_info returned JSON content with text/plain Content-Type, breaking clients that check Content-Type headers. 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/openai/chat.rs`:
- Around line 191-220: The code currently unconditionally sets CONTENT_TYPE to
"text/event-stream" for streaming paths; change this so we preserve the upstream
Content-Type for non-success responses and only force "text/event-stream" for
successful SSE responses: in the is_streaming branch (where resp.bytes_stream()
is used and you build Response::new(...)), check status.is_success() (or check
resp.headers().get(CONTENT_TYPE)) before calling
response.headers_mut().insert(CONTENT_TYPE,
HeaderValue::from_static("text/event-stream")); only insert/override the header
when the upstream status is success (or when no Content-Type exists) so 4xx/5xx
JSON error payloads keep their original Content-Type.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 13439dd3-6ce4-4c68-84e4-cb9d98d70bc8
📒 Files selected for processing (3)
model_gateway/src/routers/openai/chat.rsmodel_gateway/src/routers/openai/health.rsmodel_gateway/src/routers/openai/router.rs
2 participants
Summary
route_chatlogic (~250 lines) from the monolithicrouter.rs(~988 lines) intoopenai/chat.rsas a free function with aRouterContextstruct for shared infrastructurehealth_generateandget_server_info(~60 lines) intoopenai/health.rsrouter.rsbecomes a thinner dispatcher (~680 lines), withresolve_providerextracted as apub(super)free function shared bychat.rsandroute_responsesRefs: #721
What changed
model_gateway/src/routers/openai/chat.rsRouterContextstruct +route_chatfree function with full chat completion routing (worker selection, provider transform, retry loop, streaming/non-streaming)model_gateway/src/routers/openai/health.rshealth_generate+get_server_infofree functions withexternal_workershelper inlinedmodel_gateway/src/routers/openai/mod.rsmod chat;andmod health;model_gateway/src/routers/openai/router.rsRouterTraitmethods now delegate to extracted modules. Extractedresolve_provideraspub(super). Removed unusedshared_components()accessor.Why
router.rsmixed chat dispatch, health endpoints, responses orchestration, storage queries, and realtime methods in a single ~988-line file. Splitting improves navigability and sets up further extractions (responses in a follow-up PR, realtime after that).How
Uses free functions with a
RouterContextstruct rather than splittingimplblocks across files, sinceOpenAIRouterfields are private. This matches the convention used in the Anthropic router (RouterContext= shared infrastructure,RequestContext= per-request input) and Gemini router (streaming::execute(&router_ctx, req_ctx)).Pure refactor — no behavior changes.
Test plan
cargo build -p smg— clean, no warningscargo clippy -p smg --all-targets -- -D warnings— cleancargo test -p smg --lib— 409 tests passSummary by CodeRabbit