Skip to content

feat(rust): 1:1 port of OpenAI Responses API WebSockets to litellm-rust - #33849

Merged
ishaan-berri merged 10 commits into
litellm_internal_stagingfrom
litellm_responses_websocket_rust
Jul 19, 2026
Merged

feat(rust): 1:1 port of OpenAI Responses API WebSockets to litellm-rust#33849
ishaan-berri merged 10 commits into
litellm_internal_stagingfrom
litellm_responses_websocket_rust

Conversation

@devin-ai-integration

@devin-ai-integration devin-ai-integration Bot commented Jul 18, 2026

Copy link
Copy Markdown
Contributor

Relevant issues

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • My PR passes all CI/CD checks (e.g., lint, format, unit tests)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to 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-gateway running on 127.0.0.1:4001 with a gpt-5.3-codex deployment, Codex CLI 0.144.6 configured with a custom provider pointing at the gateway

# ~/.codex/config.toml
model = "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 = true
$ RUST_LOG=codex_api=info codex exec --skip-git-repo-check "Reply with exactly: WS-OK"
INFO responses_websocket.connect{transport="responses_websocket" api.path="responses"}: codex_api::endpoint::responses_websocket: successfully connected to websocket: ws://127.0.0.1:4001/v1/responses, headers: {"connection": "upgrade", "upgrade": "websocket", "sec-websocket-accept": "rQDhsvFHhrKPRXbG4KPqb7WJZTc=", ...}
WS-OK

A 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-gateway Axum server (like realtime), and through litellm/rust_bridge so the Python interface can run the upstream websocket over Rust behind the same rust: true deployment flag used by the /messages port. Non-native providers keep using the Python ManagedResponsesWebSocketHandler

The Rust files mirror the Python responsibilities directly:

Rust mirrors Python
core/src/responses/{types,websocket}.rs litellm/llms/base_llm/responses/transformation.py websocket surface (supports_native_websocket, get_websocket_url, model_in_websocket_url) plus typed events
core/src/providers/openai/responses/transformation.rs litellm/llms/openai/responses/transformation.py
ai-gateway/src/io/responses_ws.rs (async_responses_websocket, ResponsesWebSocketStreaming::bidirectional_forward) llm_http_handler.async_responses_websocket + litellm/responses/streaming_iterator.py
ai-gateway/src/routes/responses/ litellm/proxy/response_api_endpoints/endpoints.py responses_websocket_endpoint

Behavior ported 1:1: model from ?model= query param or extracted from the first response.create frame (flat or nested shape) with the frame replayed upstream; connection-authorized model enforced in both frame shapes; URL built as wss://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 events

Call-hook instrumentation lives entirely in core: core/src/responses/instrumentation.rs owns event accumulation (usage/model/response id), phase timing via CallLifecycle, 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_URL callback POSTs), so websocket sessions land in the proxy spend logs and the Admin UI logs page. This layering is now codified as a rule in litellm-rust/CLAUDE.md

Rust bridge (Python interface uses the Rust code, mirroring /messages):

  • crates/python-bridge: ResponsesWebSocketConnection PyO3 class (connect, send_text, recv_text, close) backed by the gateway upstream dialer
  • litellm/rust_bridge/responses_websocket.py: thin wrapper, returns None when the native module is unavailable, dependency-injectable
  • BaseLLMHTTPHandler.async_responses_websocket: when provider is openai and the deployment sets rust: true, the upstream connection is the Rust bridge object; the adapter raises ConnectionClosedOK on close so ResponsesWebSocketStreaming behaves identically; any bridge failure falls back to the Python websockets.connect path. Off by default

Tests: 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 for rust flag gating, bridge path, clean-close translation, and bridge-unavailable fallback

Final Attestation

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

Link to Devin session: https://app.devin.ai/sessions/165772ec80ba4fcfaa1c15a25b236ee4
Requested by: @ishaan-berri

devin-ai-integration Bot and others added 3 commits July 18, 2026 20:56
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>
@ishaan-berri ishaan-berri self-assigned this Jul 18, 2026
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

  • Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
  • Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

  • Disable automatic comment, CI, and merge conflict monitoring

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

@codecov

codecov Bot commented Jul 18, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 81.25000% with 15 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
litellm/rust_bridge/responses_websocket.py 86.53% 7 Missing ⚠️
litellm/llms/custom_httpx/llm_http_handler.py 57.14% 6 Missing ⚠️
litellm/rust_bridge/ocr.py 85.71% 2 Missing ⚠️

📢 Thoughts on this report? Let us know!

@greptile-apps

greptile-apps Bot commented Jul 18, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds a Rust-native WebSocket implementation of the OpenAI Responses API to the litellm-ai-gateway Axum server, porting the Python ManagedResponsesWebSocketHandler path to Rust for OpenAI-native deployments. The implementation mirrors the existing realtime route structure with URL construction, model enforcement, a splice loop with idle timeout, and a session-level usage collector.

  • Core (litellm-core): Adds typed event/error envelopes (ResponsesWsEvent, ResponsesErrorFrame), a ResponsesWebSocketProviderConfig trait with URL-building and transform hooks, and an OpenAIResponsesWsConfig singleton with passthrough transforms plus model enforcement — all unit-tested for URL parity and serialization shape.
  • Gateway (litellm-ai-gateway): Adds GET /v1/responses and GET /responses WebSocket routes behind RequireMasterKey, a tokio::select! splice loop that enforces model on client frames and observes only upstream frames for usage logging, and a ResponsesWsStreaming collector that flushes one success/failure callback when the session ends.
  • Tests: Cover key paths including flat/nested response.create, idle timeout, HTTP 401/500 dial-status mapping, pre-call error frame shape, anti-spoof observation, and pre-upgrade auth rejection.

Confidence Score: 4/5

Safe 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.

Important Files Changed

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)

  1. litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs, line 591-606 (link)

    P2 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.create event that lacks a model field, 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.

  2. litellm-rust/crates/core/src/responses/websocket.rs, line 889-898 (link)

    P2 is_terminal_event is exported but never called, and misclassifies ResponseCreated

    The function is declared pub but is never imported or invoked anywhere in the codebase; streaming.rs independently duplicates the same event-match logic. More importantly, ResponsesWsEventType::ResponseCreated is listed in the match — but that event fires at the start of a response, not at its end. The true terminal events are ResponseCompleted, ResponseFailed, ResponseIncomplete, and Error. 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

Comment on lines +27 to +32
.strip_prefix("openai/")
.unwrap_or(&params.model);
if params.model.contains('/') && !params.model.starts_with("openai/") {
return Err(CoreError::InvalidProvider(
"Responses WebSocket route supports OpenAI deployments only".to_string(),
));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 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.

devin-ai-integration Bot and others added 3 commits July 18, 2026 21:21
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-integration

Copy link
Copy Markdown
Contributor Author

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:

  • Codex CLI connects over the Rust websocket route: passed. Codex logs "successfully connected to websocket: ws://127.0.0.1:4001/v1/responses" and turn 1 created ws_proof.txt with exactly hello-from-rust-responses-websocket
  • Second response.create on the same connection: passed. Turn 2 appended second-turn-ok; a live ss watch showed the same TCP socket (codex pid to litellm-ai-gateway pid) persisting across both turns
  • Wrong master key rejected pre-upgrade: passed. codex exec with a bad key failed with HTTP error: 401 Unauthorized, url: ws://127.0.0.1:4001/v1/responses

Turn 2 on the same persistent connection (left: Codex TUI, right: live socket to :4001):

turn2

Wrong-key 401 rejection

401

Codex provider config used
model = "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 = true

Caveats: 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

@codspeed-hq

codspeed-hq Bot commented Jul 18, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_responses_websocket_rust (39a0d0d) with litellm_internal_staging (595e724)1

Open in CodSpeed

Footnotes

  1. No successful run was found on litellm_internal_staging (9533258) during the generation of this report, so 595e724 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report.

devin-ai-integration Bot and others added 2 commits July 18, 2026 22:11
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-integration

Copy link
Copy Markdown
Contributor Author

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:

  • Multi-step coding task over the Rust websocket: passed. Codex wrote and executed primes_report.py; the JSON output has the correct first 20 primes (sum 639, mean 31.95), with the live TCP socket to :4001 visible throughout
  • Second turn on the same persistent connection: passed. Codex patched the script to add largest_gap and reran it (correct value 6); same socket across both turns
  • Session lands on the Admin UI Logs page: passed. After the websocket closed, the proxy received the gateway's callback POSTs to /v1/rust_control_plane/logs and a new row appeared at /ui/logs: Success, call type responses_websocket, model gpt-5.3-codex, 64,514 tokens, duration 91.34 s, request id resp_0c83fa870989a50d016a5c077697b08199bc576ce925c3ae81, matching the LiteLLM_SpendLogs row in Postgres

Logs page row produced by the websocket session:

logs

Row detail: call type responses_websocket, tokens, duration

detail

Codex two-turn task over the websocket

turn1

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

devin-ai-integration Bot and others added 2 commits July 19, 2026 01:32
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
@ishaan-berri
ishaan-berri enabled auto-merge (squash) July 19, 2026 01:46
@ishaan-berri
ishaan-berri merged commit 7891388 into litellm_internal_staging Jul 19, 2026
76 of 77 checks passed
@ishaan-berri
ishaan-berri deleted the litellm_responses_websocket_rust branch July 19, 2026 01:55
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants