Litellm websocket improvements

[Sameerlite]

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 unit tests on make test-unit
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Screenshots / Proof of Fix

Type

🆕 New Feature
🐛 Bug Fix

Changes

---

Note

Medium Risk
Touches proxy auth (model allowlist after first frame), WebSocket request handling, and spend/router accounting; behavior changes for clients without ?model= but is covered by new tests.

Overview
Aligns the proxy Responses API WebSocket flow with OpenAI-style clients that omit ?model= and send the model on the first response.create frame (e.g. Codex).

Model resolution and first frame: ?model= is now optional. When it is missing, the server reads the first frame (30s timeout), validates JSON and type: response.create, extracts model from flat or nested response payloads, and returns structured invalid_request_error frames instead of silently dropping bad first messages.

Auth and plumbing: After the model is known from that frame, _enforce_responses_ws_first_frame_model_auth runs key allowlist and centralized checks. The consumed frame is passed as first_message through the HTTP handler, native proxy streaming, and managed WS handler so it is forwarded/processed once and not read again from the socket.

Managed WS routing: custom_llm_provider is only forced when the per-event model resolves to the same provider as the connection model (via _same_provider), so cross-provider overrides in a frame are not pinned to the wrong backend.

Spend / router: Session wrappers _aresponses_websocket and _arealtime are skipped in proxy cost tracking, router TPM/RPM updates, and budget limiter when there is no standard_logging_object, since per-turn costs are logged on inner calls.

Router alias fix: _ageneric_api_call_with_fallbacks_helper sets model to the deployment’s real model name so router aliases do not overwrite the resolved deployment model in kwargs.

^{Reviewed by Cursor Bugbot for commit 348853e. Bugbot is set up for automated code reviews on this repo. Configure here.}

[codecov]

Codecov Report

❌ Patch coverage is 89.13043% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/proxy/response_api_endpoints/endpoints.py	91.37%	5 Missing ⚠️
litellm/responses/streaming_iterator.py	86.20%	4 Missing ⚠️
litellm/router.py	50.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR aligns the Responses API WebSocket endpoint with OpenAI-style clients (e.g. Codex) that omit ?model= at connection time and send the model inside the first response.create frame. user_api_key_auth_websocket already runs full auth for URL-supplied models; a new _enforce_responses_ws_first_frame_model_auth call mirrors those checks for first-frame models, keeping coverage symmetric.

• ?model= is now optional; _read_ws_model_from_first_frame reads, validates, and extracts the model with a 30-second timeout and structured invalid_request_error frames on failure.
• _same_provider gates custom_llm_provider injection per event — improves same-provider variant routing while leaving cross-provider overrides to litellm's own resolution.
• WS session wrappers (_aresponses_websocket, _arealtime) are skipped in all three cost-tracking hooks to avoid duplicate spend accounting; per-turn costs are recorded on inner calls.
• _ageneric_api_call_with_fallbacks_helper now places \"model\": model_name last in response_kwargs, so an alias leaking through **kwargs can no longer overwrite the deployment's resolved model.

Confidence Score: 4/5

Safe to merge with attention to the double-auth-check on the first-frame path and the router-alias gap in provider injection.

The core WebSocket auth flow is correct: URL-model connections get model-allowlist checks through the existing user_api_key_auth_websocket path, and first-frame connections get equivalent checks via _enforce_responses_ws_first_frame_model_auth. The cost-tracking and router-alias fixes are straightforward and well-tested. The _same_provider logic is a genuine improvement over the previous all-or-nothing guard, but router aliases sent as per-event models are treated as unresolvable (same as before). The first-frame auth function also double-invokes _enforce_key_and_fallback_model_access and _run_centralized_common_checks relative to common_processing_pre_call_logic, adding latency on first-frame connections.

litellm/proxy/response_api_endpoints/endpoints.py (double auth call on first-frame path) and litellm/responses/streaming_iterator.py (_same_provider alias behaviour)

Important Files Changed

Filename	Overview
litellm/proxy/response_api_endpoints/endpoints.py	Makes ?model= optional; adds _read_ws_model_from_first_frame, _extract_model_from_first_ws_event, and _enforce_responses_ws_first_frame_model_auth. Model-specific auth (allowlist + centralized checks) only runs when model comes from the first frame, not from the URL query param — an intentional asymmetry that warrants attention.
litellm/responses/streaming_iterator.py	Adds first_message forwarding through both streaming paths; introduces _same_provider and _connection_provider to gate custom_llm_provider injection. The _same_provider logic returns False for unresolvable event models (e.g. router aliases), preventing provider inheritance that the old code also didn't supply — net behavior is the same for that case.
litellm/router.py	Two changes: explicit "model": model_name at end of response_kwargs so alias from kwargs can't overwrite deployment model; early return in _update_tpm_rpm_redis_async for WS session wrapper call types. Both changes are correct and well-targeted.
litellm/proxy/hooks/proxy_track_cost_callback.py	Adds skip guard for _aresponses_websocket and _arealtime call types when standard_logging_object is None; per-turn costs are tracked on inner calls. Simple and correct.
litellm/router_strategy/budget_limiter.py	Same skip guard as proxy_track_cost_callback.py for WS session wrappers. Simple one-liner, correctly placed before payload extraction.
litellm/llms/custom_httpx/llm_http_handler.py	Threads first_message through aresponses_websocket_handler and ResponsesWebSocketStreaming constructors; plumbing-only change, no logic changes.
tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py	Adds extensive mock-only tests for first-frame model extraction, validation, auth enforcement, cost-tracking skips, and _same_provider logic. All tests use mocks with no real network calls. Coverage is thorough for the happy paths.
tests/test_litellm/test_router.py	Adds test_ageneric_api_call_deployment_model_overrides_alias with inject_alias_into_kwargs side_effect to replicate the real call-path alias leak — an improvement over what was flagged in the previous review thread.

_{Reviews (12): Last reviewed commit: "fix(responses-ws): fall back to explicit..." | Re-trigger Greptile}

[veria-ai]

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 1 · PR risk: 0/10

[Sameerlite]

@greptileai

[Sameerlite]

@greptileai

[Sameerlite]

@greptileai

[Sameerlite]

@greptileai

[Sameerlite]

@greptileai

[Sameerlite]

@greptileai

[Sameerlite]

bugbot run

[cursor]

Cursor Bugbot has reviewed your changes using high effort and found 3 potential issues.

Autofix Details

Bugbot Autofix prepared fixes for all 3 issues found in the latest run.

• ✅ Fixed: Consumed first frame dropped silently
  • The first frame is now rejected with a WebSocket error unless it is a response.create JSON object, so it is never silently consumed.
• ✅ Fixed: Auth skips model allowlist
  • When the model is resolved from the first frame, the proxy now reruns key and common model authorization checks before routing.
• ✅ Fixed: Non-object JSON crashes handler
  • The first-frame parser now validates the decoded payload type and handles non-object JSON as a structured invalid request.

_{You can send follow-ups to the cloud agent here.}

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
2 out of 3 committers have signed the CLA.

✅ Sameerlite
✅ mateo-berri
❌ cursoragent
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[Sameerlite]

bugbot run

[cursor]

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

^{Reviewed by Cursor Bugbot for commit 348853e. Configure here.}

[Sameerlite]

@greptileai

[mateo-berri]

@greptileai

[mateo-berri]

@greptileai

[mateo-berri]

@greptileai

[mateo-berri]

@greptileai

[mateo-berri]

LGTM; thanks!

[krrish-berri-2]

i'm confused - does this add wss:// /v1/responses to litellm? codex is complaining this endpoint doesn't exist

[krrish-berri-2]

this is causing a 4.5s startup latency when launching codex for a session via litellm

[Sameerlite]

i'm confused - does this add wss:// /v1/responses to litellm? codex is complaining this endpoint doesn't exist

So this was designed based on how we have designed realtime websocket. Model should be present at the time connction as query param and then based on that we will connect it to the correct provider. But openai was first connecting. And then sending model in codex. Hence it was not working. I updated the code to accept the websocket req first and when the model is sent, use that for routing and starting actual websocket connection with the provider

this is causing a 4.5s startup latency when launching codex for a session via litellm

Not able to reproduce
https://www.loom.com/share/6541b419e87d43ba815d01c67d5c39c9