Repository navigation
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces WebSocket support for the Responses API, which is a significant feature enhancement. The implementation is comprehensive, including a new CLI argument for connection limits, state management, a new WebSocket handler module, and extensive tests (unit, integration, and e2e). The code is generally well-structured and robust. I've found one critical issue in the WebSocket connection handling logic that could lead to a runtime error, and a related redundancy. My review includes suggestions to fix this. Overall, this is a great contribution.
| app = websocket.app | ||
| serving = app.state.openai_serving_responses | ||
|
|
||
| if serving is None: | ||
| await websocket.close(code=1008, | ||
| reason="Responses API not supported") | ||
| return | ||
|
|
||
| # Connection limit check | ||
| async with app.state.ws_responses_lock: | ||
| if (app.state.ws_responses_active_connections | ||
| >= app.state.ws_responses_max_connections): | ||
| await websocket.accept() | ||
| await websocket.send_text(json.dumps({ | ||
| "type": "error", | ||
| "status": 429, | ||
| "error": { | ||
| "code": "websocket_connection_limit_reached", | ||
| "message": "Maximum WebSocket connections reached", | ||
| }, | ||
| })) | ||
| await websocket.close() | ||
| return | ||
| app.state.ws_responses_active_connections += 1 |
There was a problem hiding this comment.
The current WebSocket connection handling logic has a bug and can be simplified.
- Bug: If
serving is None,websocket.close()is called beforewebsocket.accept(), which will raise aRuntimeError: Cannot close a WebSocket connection that is not open. - Redundancy:
websocket.accept()is called in the connection limit check path, and also insideWebSocketResponsesConnection.handle_connection(). This is redundant.
To fix this, you should accept the connection once at the beginning of the create_responses_websocket function. This ensures the connection is open before any other operations and simplifies the logic.
await websocket.accept()
app = websocket.app
serving = app.state.openai_serving_responses
if serving is None:
await websocket.close(code=1008,
reason="Responses API not supported")
return
# Connection limit check
async with app.state.ws_responses_lock:
if (app.state.ws_responses_active_connections
>= app.state.ws_responses_max_connections):
await websocket.send_text(json.dumps({
"type": "error",
"status": 429,
"error": {
"code": "websocket_connection_limit_reached",
"message": "Maximum WebSocket connections reached",
},
}))
await websocket.close()
return
app.state.ws_responses_active_connections += 1There was a problem hiding this comment.
Thanks for catching this! Both issues are already resolved in the current code:
- Bug fixed:
websocket.accept()is now called first (line 155), before theserving is Nonecheck — so we never attempt to close an unaccepted connection. - No redundancy:
websocket.py(handle_connection) does not callaccept()— it's only called once in the router.
The current flow matches exactly what you suggested.
| """Main WebSocket connection loop.""" | ||
| from starlette.websockets import WebSocketDisconnect | ||
|
|
||
| await self.websocket.accept() |
There was a problem hiding this comment.
Already addressed — handle_connection() in websocket.py does not call accept(). The connection is accepted once in the router before handle_connection is called, so there's no redundancy.
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run You ask your reviewers to trigger select CI tests on top of Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. 🚀 |
437831b to
2128b46
Compare
| @@ -137,5 +137,55 @@ async def cancel_responses(response_id: str, raw_request: Request): | |||
| return JSONResponse(content=response.model_dump()) | |||
|
|
|||
|
|
|||
| @router.websocket("/v1/responses") | |||
There was a problem hiding this comment.
Just a quick question: is this compatible with the OpenAI Responses API WebSocket standard?
Could you also write some E2E tests?
There was a problem hiding this comment.
Thanks for the review @chaunceyjiang!
OpenAI compatibility — yes, this implementation follows the OpenAI Responses API WebSocket standard. Specifically:
| Spec requirement | Implementation |
|---|---|
Endpoint ws(s)://host/v1/responses |
@router.websocket("/v1/responses") in api_router.py |
Client sends {"type": "response.create", ...} |
handle_event() dispatches on type == "response.create" |
| Server streams same events as SSE mode | _process_response_create() forwards the existing streaming pipeline events over WebSocket frames |
stream / background fields omitted from WS payload |
Stripped in _process_response_create() (payload["stream"] = True, payload.pop("background", None)) |
| One in-flight request per connection | ctx.inflight guard with concurrent_request error (409) |
previous_response_id for continuation via connection-local cache |
ConnectionContext.last_response_id/last_response with previous_response_not_found error |
generate: false warmup mode |
Returns response.created + response.completed with no inference |
| 60-minute connection lifetime | _enforce_deadline() background task (warning at 55 min, close at 60 min) |
Error format {type: "error", status, error: {code, message}} |
send_error() method |
E2E tests — already included in this PR at tests/entrypoints/openai/responses/test_websocket_real.py. These spin up a real vLLM server with a tiny model and exercise the full WebSocket flow end-to-end. Coverage includes:
- Basic streaming (connect →
response.create→ receive events →response.completed) - Streaming event ordering (created → in_progress → output_item.added → deltas → done → completed)
- Text consistency (concatenated deltas == final output text)
- Warmup mode (
generate: false) - Continuation with
previous_response_id - Invalid
previous_response_iderror - Unknown event type error (connection stays open)
- Invalid JSON error (connection stays open)
- Multiple sequential requests on same connection
- Multiple concurrent connections
- Warmup → real request continuation
- Response usage/token counts
In addition, there are unit tests (test_websocket.py) and integration tests with mock serving (test_websocket_integration.py).
There was a problem hiding this comment.
@rasonyang Thanks~
I have another question: when using --api-server-count=2, does WebSocket work properly?
There was a problem hiding this comment.
Good question. WebSocket works with --api-server-count=2, but with two caveats:
--max-websocket-connectionsis per-process, not global. With 2 servers and limit=100, the actual max is 200.- Connection-local state (
previous_response_idcache) is per-process. This is fine in practice — a single WebSocket connection is a persistent TCP connection pinned to one process for its entire lifetime, so the cache stays consistent. The state only matters within the same connection, and connections don't migrate between processes.
So it works correctly — the only nuance is the per-process connection limit. I can add a note in the docstring if you'd like.
|
Hi @rasonyang, the pre-commit checks have failed. Please run: uv pip install pre-commit
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, Tip Is
|
b938069 to
399a6e4
Compare
|
Hi @rasonyang, the pre-commit checks have failed. Please run: uv pip install pre-commit
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, Tip Is
|
fc8757d to
a058383
Compare
| json.dumps( | ||
| { | ||
| "type": "error", | ||
| "status": 429, |
There was a problem hiding this comment.
Could we use a constant instead of a magic number?
There was a problem hiding this comment.
Good catch — replaced 429 with HTTPStatus.TOO_MANY_REQUESTS. Fixed in the latest push.
a058383 to
e6d4125
Compare
|
Hi @rasonyang, the pre-commit checks have failed. Please run: uv pip install pre-commit
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, Tip Is
|
65c798d to
4ea86ae
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
b769df5 to
12e1802
Compare
12e1802 to
a272270
Compare
|
This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you! |
d3ee956 to
7a96299
Compare
70dd9c1 to
969c2a8
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
Add WebSocket transport for the OpenAI Responses API, enabling persistent bidirectional connections for lower-latency streaming. Changes: - Add `--max-websocket-connections` CLI argument (default: 100) - Initialize WebSocket server state (semaphore, active connections counter) in `init_app_state()` - Add WebSocket route `/v1/responses/websocket` with connection limit enforcement in the responses API router - New `websocket.py` module with full connection lifecycle: ConnectionContext, uuid7-based event IDs, JSON message parsing, streaming response events, warmup/error handling, and graceful shutdown Tests: - Unit tests (test_websocket.py): ConnectionContext, uuid7, message parsing, event formatting — no GPU/torch required - Mock integration tests (test_websocket_integration.py): full connection lifecycle using Starlette TestClient with mocked engine - Real e2e tests (test_websocket_real.py): live CPU server with facebook/opt-125m, covering streaming, multi-turn, cancellation, and concurrent connections Signed-off-by: Rason Yang <rason.yang@gmail.com>
Add real GPU integration tests for the WebSocket mode of the Responses API. These tests run against a live vLLM server with GPU inference using the Qwen/Qwen2-0.5B-Instruct model. The test suite covers: - Basic streaming - Event ordering - Text consistency - Warmup requests - Response continuation with previous_response_id - Error handling (invalid JSON, unknown events) - Sequential and concurrent requests - Usage statistics This is the GPU counterpart to test_websocket_real.py (CPU version). Signed-off-by: Rason Yang <rason.yang@gmail.com>
…GPU tests - Update shell script and docstring references from facebook/opt-125m to Qwen/Qwen2-0.5B-Instruct to match the actual MODEL_NAME constant - Rename _ws to ws since the variable is actively used Signed-off-by: Rason Yang <rason.yang@gmail.com>
…ff E501 Signed-off-by: Rason Yang <rason.yang@gmail.com>
Upstream moved `load_aware_call`, `validate_json_request` and `with_cancellation` from `vllm.entrypoints.utils` / `vllm.entrypoints.openai.utils` into `vllm.entrypoints.serve.utils.api_utils`. The torch-free import stubs in the WebSocket tests still registered the old module paths, so importing `responses/api_router.py` no longer resolved through them. Point the stubs at the new module. Signed-off-by: rason yang <rason.yang@gmail.com>
…lient Raw exception text was forwarded to remote WebSocket clients from both `processing_error` paths, which can leak server file paths, memory addresses and tracebacks. Route both through `sanitize_message()`, matching the equivalent handling in `speech_to_text/realtime/connection.py`. The unsanitized detail is still recorded server-side via `logger.exception`. Errors arriving as an `ErrorResponse` from the serving layer are already sanitized by `create_error_response`, so they are unchanged. Importing `sanitize_message` executes `vllm.entrypoints.serve.__init__`, which pulls `ErrorInfo` and `GenerationError` from the engine protocol module, so extend the torch-free import stubs in the tests accordingly. Add a regression test asserting paths and memory addresses are stripped while the underlying cause stays visible. Signed-off-by: rason yang <rason.yang@gmail.com>
…rocess The limit is enforced independently in each API server process, so with --api-server-count=N the effective total is N times the configured value. Raised by @chaunceyjiang in review. Signed-off-by: rason yang <rason.yang@gmail.com>
…orts Upstream vllm-project#53659 moved cli_args.py from vllm/entrypoints/openai/ to vllm/entrypoints/launchers/, which broke the AST check that reads the file by path. Also install the module stubs at import time. websocket.py imports sanitize_message, which executes vllm/entrypoints/serve/__init__.py and pulls in torch, so the tests above only ran on torch-less machines by accident of an earlier failed import leaving sys.modules populated. Guarded on torch being absent, so CI behaviour is unchanged. Signed-off-by: rason yang <rason.yang@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
caec653 to
d8ce8bd
Compare
|
Rebased onto current Upstream #53659 moved Addressed since the last review pass
The earlier questions on OpenAI WebSocket-mode spec compatibility and Testing
CI
That skips I have reviewed every changed line in this PR. AI assistance was used, as noted in the description. |
|
@chaunceyjiang could you add the |
|
we should coordinate this with Agentic API before adding another WebSocket continuation implementation. Agentic API already handles Responses WebSockets, connection-local state, prewarming, persistence, and tool orchestration. its upstream path uses stateless the non-stateful Responses scope in #44280 and the work in #53380 look like the right backend boundary for that integration. standalone there are also a few correctness issues in
|
|
Thanks @franciscojavierarceo. I verified all four issues against Agreed on coordinating with Agentic API. Since it only needs stateless Is there a thread tracking the Python / Rust frontend / Agentic API split? |
|
I raised a discussion in the sig-frontend slack channel here |
Summary
Add WebSocket transport for the OpenAI Responses API, enabling persistent bidirectional connections for lower-latency streaming.
--max-websocket-connectionsCLI argument (default: 100, enforced per API server process) to control server capacityinit_app_state()/v1/responseswith connection limit enforcementwebsocket.pymodule:ConnectionContext, uuid7-based event IDs, JSON message parsing, streaming response events, warmup/error handling, and graceful shutdownFollows the OpenAI Responses API WebSocket protocol: client connects to
ws(s)://host/v1/responses, sends{"type": "response.create", ...}, and the server streams back the same event objects as SSE mode.Files Changed
vllm/entrypoints/launchers/cli_args.py--max-websocket-connectionsargvllm/entrypoints/launchers/api_server/app_state.pyinit_app_state()vllm/entrypoints/openai/responses/api_router.pyvllm/entrypoints/openai/responses/websocket.pytests/.../test_websocket.pytests/.../test_websocket_integration.pytests/.../test_websocket_real.pyQwen/Qwen2-0.5B-Instruct)tests/.../test_websocket_gpu.pyQwen/Qwen2-0.5B-Instruct)scripts/run_websocket_gpu_tests.shNotes
--api-server-count=N, the effective total isN * --max-websocket-connections. Documented in the CLI arg docstring.previous_response_idcontinuation is connection-local. It resolves against a per-connection cache rather than the shared response store. A WebSocket connection is pinned to one process for its lifetime, so this stays consistent under multiple API servers.init_app_state()now lives invllm/entrypoints/launchers/api_server/app_state.py— upstream turnedvllm/entrypoints/openai/api_server.pyinto a deprecation shim, so this PR no longer touches that file.cli_args.pyfromvllm/entrypoints/openai/tovllm/entrypoints/launchers/, so--max-websocket-connectionsnow lands in the new path.Test Plan
Results below are from a full run before the rebase onto current
main(commitd3ee956b0):ConnectionContext,uuid7, message parsing, event formattingQwen/Qwen2-0.5B-Instructcovering streaming, multi-turn, cancellation, concurrent connections, and error handlingQwen/Qwen2-0.5B-Instruct, same coverageRe-run after the rebase onto current
main:ruff check/ruff format --check(v0.14.0, matching.pre-commit-config.yaml) on all changed files — cleantests/entrypoints/openai/responses/conftest.pyimportstests.utils(which needs torch), so they were executed by importing the modules directly rather than through pytest.test_websocket_real.pyandtest_websocket_gpu.py— not re-run post-rebase; they need a server and a GPU box.pre-run-checkfails withPR must have the 'verified', 'ready', or 'ready-run-all-tests' label to run pre-commit, or the author must have at least 4 merged PRs (found 0), which skipspre-commitand gates Buildkite off entirely. This PR has therefore never had a CI run. A maintainer applyingreadywould let CI validate it; the workflow triggers onlabeled, so no new push is needed.Model evaluation
Not applicable — this adds a transport for existing Responses API functionality. Generation goes through the unchanged
create_responses()serving pipeline, so model output and accuracy are unaffected. No sampling, tokenization, or scheduling code is touched.Duplicate-work check
gh pr list --repo vllm-project/vllm --state open --search "websocket responses"returns only this PR. The one adjacent open PR, #38084 (/v1/responses/renderendpoint), is a separate feature and does not overlap.AI assistance
AI assistance was used for parts of this change, including the rebases onto current
main, the conflict resolution, thesanitize_messagefix for the error paths flagged in review, and the test fixes that followed upstream's module relocations.