fix(realtime-api): worker health tracking in websocket session - #725
Conversation
Signed-off-by: yifeliu <yifengliu9@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
📝 WalkthroughWalkthroughModified the OpenAI realtime WebSocket handler to replace direct error handling with a guard-based approach that returns a success indicator, enabling outcome recording via Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related issues
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 |
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 addresses a critical gap in worker health tracking for real-time WebSocket sessions. Previously, the system lacked a mechanism to report the outcome of these sessions, preventing the load balancer from accurately assessing worker health. The changes introduce a robust way to capture the success or failure of WebSocket proxy operations and record these outcomes, bringing the real-time path in line with existing patterns in the REST API and ensuring more reliable load balancing. 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
The pull request refactors the real-time WebSocket proxy handling to explicitly capture and record the success or failure of the proxy operation. This involves updating tracing imports, adjusting the proxy module import, and introducing a worker.record_outcome call to track the session's outcome before cleanup.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
model_gateway/src/routers/openai/realtime/ws.rs (2)
120-137:⚠️ Potential issue | 🟠 Major
successstill over-reports healthy WS sessions.Lines 120-137 treat any
Ok(())fromproxy::run_ws_proxyas success, butmodel_gateway/src/routers/openai/realtime/proxy.rscurrently returnsOk(())for all post-connect exits and only logs task failures before falling through. That means a session can terminate abnormally after connect and still callworker.record_outcome(true), so the health signal remains misleading. Please propagate a richer proxy outcome here, or make abnormal forwarding/session termination returnErrinstead ofOk(()).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/openai/realtime/ws.rs` around lines 120 - 137, The code treats any Ok(()) from proxy::run_ws_proxy as a healthy session, but proxy::run_ws_proxy currently returns Ok(()) for abnormal post-connect exits, causing worker.record_outcome(true) to over-report health; change the proxy API or caller to propagate a richer outcome and only count true for genuinely successful sessions. Specifically, update proxy::run_ws_proxy to return a Result<ProxyOutcome, E> (or Result<bool, E>) that distinguishes normal completion vs abnormal termination, modify the match here to treat abnormal outcomes as Err/false (e.g., match proxy::run_ws_proxy(...) .await { Ok(ProxyOutcome::Completed) => true, Ok(ProxyOutcome::Aborted) | Err(_) => { error!(...); false } }), and then pass that boolean into worker.record_outcome so only true reflects a truly healthy forwarded session; ensure references to proxy::run_ws_proxy and worker.record_outcome are updated accordingly.
119-139: 🧹 Nitpick | 🔵 TrivialAdd a regression test for the new health accounting.
This branch changes worker health state, but there is no coverage here to lock down the
true/falsemapping. A focused test with a fakeWorkerthat assertsrecord_outcome(true)on a clean proxy completion andrecord_outcome(false)on proxy failure would make this much safer to maintain.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/openai/realtime/ws.rs` around lines 119 - 139, Add a regression test that asserts the worker health accounting mapping by creating a fake/mock Worker and exercising the ws upgrade path: call the code that invokes proxy::run_ws_proxy (or directly invoke the closure used in ws.on_upgrade) with two scenarios—one where proxy::run_ws_proxy returns Ok(()) and one where it returns Err(...). For each scenario, verify FakeWorker.record_outcome(true) is called on the clean completion case and FakeWorker.record_outcome(false) is called on the failure case; use the same symbols from the diff (proxy::run_ws_proxy, Worker.record_outcome, realtime_registry.remove_session, session_id) to locate and wire the test, and inject/mocking run_ws_proxy or the WebSocket input so the test controls success vs failure deterministically. Ensure the test runs under tokio async context and cleans up realtime_registry entries as the production code does.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@model_gateway/src/routers/openai/realtime/ws.rs`:
- Around line 120-137: The code treats any Ok(()) from proxy::run_ws_proxy as a
healthy session, but proxy::run_ws_proxy currently returns Ok(()) for abnormal
post-connect exits, causing worker.record_outcome(true) to over-report health;
change the proxy API or caller to propagate a richer outcome and only count true
for genuinely successful sessions. Specifically, update proxy::run_ws_proxy to
return a Result<ProxyOutcome, E> (or Result<bool, E>) that distinguishes normal
completion vs abnormal termination, modify the match here to treat abnormal
outcomes as Err/false (e.g., match proxy::run_ws_proxy(...) .await {
Ok(ProxyOutcome::Completed) => true, Ok(ProxyOutcome::Aborted) | Err(_) => {
error!(...); false } }), and then pass that boolean into worker.record_outcome
so only true reflects a truly healthy forwarded session; ensure references to
proxy::run_ws_proxy and worker.record_outcome are updated accordingly.
- Around line 119-139: Add a regression test that asserts the worker health
accounting mapping by creating a fake/mock Worker and exercising the ws upgrade
path: call the code that invokes proxy::run_ws_proxy (or directly invoke the
closure used in ws.on_upgrade) with two scenarios—one where proxy::run_ws_proxy
returns Ok(()) and one where it returns Err(...). For each scenario, verify
FakeWorker.record_outcome(true) is called on the clean completion case and
FakeWorker.record_outcome(false) is called on the failure case; use the same
symbols from the diff (proxy::run_ws_proxy, Worker.record_outcome,
realtime_registry.remove_session, session_id) to locate and wire the test, and
inject/mocking run_ws_proxy or the WebSocket input so the test controls success
vs failure deterministically. Ensure the test runs under tokio async context and
cleans up realtime_registry entries as the production code does.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: e5b47a4a-bd70-4c4d-8603-3a6c6916afac
📒 Files selected for processing (1)
model_gateway/src/routers/openai/realtime/ws.rs
2 participants
Ref: #637
Description
Problem
The realtime WebSocket handler does not call
worker.record_outcome()after the proxy session completes. This means worker health tracking has no signal from realtime WebSocket sessions — unlike the REST path which already records outcomes. Without this, the load balancer cannot distinguish healthy workers from unhealthy ones based on realtime session results.Solution
Capture the proxy result as a success
booleanand callworker.record_outcome(success)before cleaning up the session. This brings the WS handler in line with existing patterns in the REST path.Changes
model_gateway/src/routers/openai/realtime/ws.rsTest Plan
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
Refactor
Chores