feat(rust): 1:1 port of OpenAI Responses API WebSockets to litellm-rust - #33849
Conversation
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
|
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR adds a Rust-native WebSocket implementation of the OpenAI Responses API to the
Confidence Score: 4/5Safe to merge with minor polish; the WebSocket proxy path is well-tested and the auth/model-enforcement logic is sound. The splice loop, model enforcement, idle timeout, and anti-spoof observation are all implemented correctly and backed by thorough tests. Three issues were found: the first-frame validation in bridge() checks model presence before event type so a wrong-type frame without a model field produces a misleading error message; the exported is_terminal_event function is never called and incorrectly lists ResponseCreated as terminal; and non-OpenAI provider misconfigurations surface as opaque 1011 closes with no explanatory error frame. None of these affect the happy path or security posture. routes/responses/mod.rs (first-frame validation ordering) and core/src/responses/websocket.rs (is_terminal_event dead code and wrong event classification) are the files most worth a second look before merge.
|
| Filename | Overview |
|---|---|
| litellm-rust/crates/ai-gateway/src/io/responses_ws.rs | New WebSocket splice loop connecting client and upstream; idle timeout, model enforcement, and HTTP-status preservation on dial failure are all implemented and tested. Auth header injection and key resolution are correct. |
| litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs | Route handler and bridge logic; pre-upgrade auth and model validation are in place, but the first-frame validation checks model presence before event type, giving misleading error messages when the event type is wrong. |
| litellm-rust/crates/ai-gateway/src/routes/responses/service.rs | Provider-type guard for OpenAI-only deployments runs after WebSocket upgrade; non-OpenAI provider errors surface as opaque 1011 closes with no descriptive error frame. |
| litellm-rust/crates/core/src/responses/websocket.rs | URL construction, model enforcement, and provider trait are well-tested and match the Python proxy shape; exported is_terminal_event is dead code and misclassifies ResponseCreated as a terminal event. |
| litellm-rust/crates/core/src/responses/types.rs | Type definitions for events and error frames; serialization round-trips and error shape are unit-tested and correct. |
| litellm-rust/crates/ai-gateway/src/responses/streaming.rs | Session-level usage accumulator and logging callback; response_cost is hardcoded to 0.0 (same pattern as realtime module), and usage accumulation with += is safe for single-terminal-event sessions. |
| litellm-rust/crates/core/src/providers/openai/responses/transformation.rs | OpenAI provider config with native WebSocket support and passthrough transforms; model enforcement tested. |
| litellm-rust/crates/core/src/constants.rs | New constants file with correct OpenAI API base and responses path values. |
Comments Outside Diff (2)
-
litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs, line 591-606 (link)Event type check occurs after model extraction, producing misleading error messages
The model is extracted (and validated for presence) before the event type is checked. If the first frame is a non-
response.createevent that lacks amodelfield, the client receives"Missing model in response.create event"rather than"First frame must be a response.create event", because the model-presence guard fires first. The event type guard should run first to give the client the most accurate diagnostic. -
litellm-rust/crates/core/src/responses/websocket.rs, line 889-898 (link)is_terminal_eventis exported but never called, and misclassifiesResponseCreatedThe function is declared
pubbut is never imported or invoked anywhere in the codebase;streaming.rsindependently duplicates the same event-match logic. More importantly,ResponsesWsEventType::ResponseCreatedis listed in the match — but that event fires at the start of a response, not at its end. The true terminal events areResponseCompleted,ResponseFailed,ResponseIncomplete, andError. If this function is ever wired up, it will incorrectly treat session-start events as session-ending ones.
Reviews (1): Last reviewed commit: "fix(rust): align Responses WebSocket par..." | Re-trigger Greptile
| .strip_prefix("openai/") | ||
| .unwrap_or(¶ms.model); | ||
| if params.model.contains('/') && !params.model.starts_with("openai/") { | ||
| return Err(CoreError::InvalidProvider( | ||
| "Responses WebSocket route supports OpenAI deployments only".to_string(), | ||
| )); |
There was a problem hiding this comment.
Non-OpenAI provider errors close 1011 with no error frame
When a configured deployment uses a non-OpenAI provider (e.g. vertex_ai/gemini-1.5-pro) the InvalidProvider error propagates to bridge(), which closes the socket with code 1011 "Internal server error" and no preceding error frame. The client cannot distinguish a misconfiguration from a real internal failure. Emitting a ResponsesErrorFrame::invalid_request(...) before the 1011 close — or checking the provider type before the WebSocket upgrade in validate_model — would make this actionable for the caller.
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
|
Live e2e test of this PR at head commit 3befd68: built the release gateway, ran it on 127.0.0.1:4001 with a real OPENAI_API_KEY and a gpt-5.3-codex deployment, and drove it end to end with Codex CLI 0.144.6 configured as a custom Responses provider with supports_websockets = true Results:
Turn 2 on the same persistent connection (left: Codex TUI, right: live socket to :4001): Codex provider config usedmodel = "gpt-5.3-codex"
model_provider = "litellm"
[model_providers.litellm]
name = "LiteLLM Rust AI Gateway"
base_url = "http://127.0.0.1:4001/v1"
wire_api = "responses"
env_key = "LITELLM_GW_KEY"
supports_websockets = trueCaveats: the gateway's logging callback POSTs to a Python proxy on :4000 were not exercised (no proxy running in this standalone setup), and the rust: true Python bridge path is covered by unit tests, not this live demo Full annotated screen recording shared in the session: https://app.devin.ai/sessions/165772ec80ba4fcfaa1c15a25b236ee4 |
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
|
Second live e2e round at head commit f6c6a4a, this time with the full logging pipeline: release gateway on 127.0.0.1:4001 with LITELLM_PROXY_BASE_URL pointing at a Python proxy on :4000 backed by Postgres, driven by Codex CLI over the Rust websocket with a real OpenAI key Results:
Logs page row produced by the websocket session: Caveats: cost displays $0.00 because gpt-5.3-codex has no local price entry (token accounting is correct), and the sole failing CI check (code-quality) is a dependency license verification failure for vcrpy and locust, unrelated since this PR adds no Python dependencies Devin session: https://app.devin.ai/sessions/165772ec80ba4fcfaa1c15a25b236ee4 |
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
7891388
into
litellm_internal_staging


Relevant issues
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Live e2e proof, no mocks, real OpenAI spend, at head commit:
litellm-ai-gatewayrunning on127.0.0.1:4001with agpt-5.3-codexdeployment, Codex CLI 0.144.6 configured with a custom provider pointing at the gatewayA screen recording of a full interactive Codex TUI session over the Rust websocket is being shared separately
Type
🆕 New Feature
Changes
1:1 port of the OpenAI-native Responses API WebSocket path to litellm-rust, exposed two ways: as a route on the
litellm-ai-gatewayAxum server (likerealtime), and throughlitellm/rust_bridgeso the Python interface can run the upstream websocket over Rust behind the samerust: truedeployment flag used by the/messagesport. Non-native providers keep using the PythonManagedResponsesWebSocketHandlerThe Rust files mirror the Python responsibilities directly:
core/src/responses/{types,websocket}.rslitellm/llms/base_llm/responses/transformation.pywebsocket surface (supports_native_websocket,get_websocket_url,model_in_websocket_url) plus typed eventscore/src/providers/openai/responses/transformation.rslitellm/llms/openai/responses/transformation.pyai-gateway/src/io/responses_ws.rs(async_responses_websocket,ResponsesWebSocketStreaming::bidirectional_forward)llm_http_handler.async_responses_websocket+litellm/responses/streaming_iterator.pyai-gateway/src/routes/responses/litellm/proxy/response_api_endpoints/endpoints.pyresponses_websocket_endpointBehavior ported 1:1: model from
?model=query param or extracted from the firstresponse.createframe (flat or nested shape) with the frame replayed upstream; connection-authorized model enforced in both frame shapes; URL built aswss://api.openai.com/v1/responses?model=<m>with the same default-base/scheme-flip/append rules as Python; pre-call failures send{"type":"error","error":{"type":"invalid_request_error","message":...}}then close 1008 "Pre-call error"; internal failures close 1011 with no detail-carrying frame; dial failures preserve the provider HTTP status; only upstream-to-client events feed logging so clients cannot spoof usage-bearing terminal eventsCall-hook instrumentation lives entirely in core:
core/src/responses/instrumentation.rsowns event accumulation (usage/model/response id), phase timing viaCallLifecycle, and success/failure callback payload construction. The gateway host only feeds observed upstream events into core and dispatches the completed payloads through its logger I/O (LiteLLM_PROXY_BASE_URLcallback POSTs), so websocket sessions land in the proxy spend logs and the Admin UI logs page. This layering is now codified as a rule inlitellm-rust/CLAUDE.mdRust bridge (Python interface uses the Rust code, mirroring
/messages):crates/python-bridge:ResponsesWebSocketConnectionPyO3 class (connect,send_text,recv_text,close) backed by the gateway upstream dialerlitellm/rust_bridge/responses_websocket.py: thin wrapper, returnsNonewhen the native module is unavailable, dependency-injectableBaseLLMHTTPHandler.async_responses_websocket: when provider isopenaiand the deployment setsrust: true, the upstream connection is the Rust bridge object; the adapter raisesConnectionClosedOKon close soResponsesWebSocketStreamingbehaves identically; any bridge failure falls back to the Pythonwebsockets.connectpath. Off by defaultTests: core unit tests for URL parity, model enforcement, event/error-frame shape; gateway tests against a mock tokio-tungstenite upstream (flat/nested
response.create, sequential requests on one connection, passthrough fidelity, anti-spoof observation, idle timeout, 401/500 dial status mapping, pre-call error frame + 1008, pre-upgrade auth rejection); Python tests forrustflag gating, bridge path, clean-close translation, and bridge-unavailable fallbackFinal Attestation
Link to Devin session: https://app.devin.ai/sessions/165772ec80ba4fcfaa1c15a25b236ee4
Requested by: @ishaan-berri