[Frontend] Consolidate entrypoint middleware (#52309)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Aug 17, 2026 at 14:17 UTC
93550cc4cd7e97bebe47b9b57333b60921c0d662
14 files changed
+567
-370
tests/entrypoints/serve/middleware/__init__.py
tests/entrypoints/serve/middleware/test_authentication_middleware.py
new
+150
@@ -0,0 +1,150 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from argparse import Namespace
5
+from typing import get_args
6
+
7
+import pytest
8
+import regex as re
9
+from fastapi import FastAPI
10
+from starlette.responses import JSONResponse
11
+from starlette.routing import Route
12
+from starlette.testclient import TestClient
13
+
14
+from vllm.entrypoints.launchers.api_server.routers import register_api_routers
15
+from vllm.entrypoints.serve.middleware.authenticate import (
16
+ GUARDED_PREFIX,
17
+ AuthenticationMiddleware,
18
+)
19
+from vllm.tasks import POOLING_TASKS, SupportedTask
20
+
21
+# ---------------------------------------------------------------------------
22
+# Helpers
23
+# ---------------------------------------------------------------------------
24
+
25
+
26
+def get_all_http_routes(app: FastAPI) -> list[tuple[str, list[str]]]:
27
+ """Extract all HTTP routes (path, methods) from the FastAPI app."""
28
+ routes = []
29
+ for route in app.routes:
30
+ if not isinstance(route, Route):
31
+ continue
32
+ path = route.path
33
+ methods = list(route.methods or {"GET"})
34
+ routes.append((path, methods))
35
+ return routes
36
+
37
+
38
+def generate_test_path(path_template: str) -> str:
39
+ """Replace path parameters (e.g. {response_id}) with 'test'."""
40
+ return re.sub(r"\{[^}]+\}", "test", path_template)
41
+
42
+
43
+def _create_app_with_mock_routes(routes: list[tuple[str, list[str]]]) -> FastAPI:
44
+ """Create a FastAPI app with AuthenticationMiddleware and mock endpoints."""
45
+ app = FastAPI()
46
+ app.add_middleware(AuthenticationMiddleware, tokens=["valid-token"])
47
+
48
+ async def mock_endpoint():
49
+ return JSONResponse({"status": "ok"})
50
+
51
+ for path_template, methods in routes:
52
+ allowed_methods = list(set(methods + ["OPTIONS"]))
53
+ app.add_api_route(
54
+ path_template,
55
+ mock_endpoint,
56
+ methods=allowed_methods,
57
+ include_in_schema=False,
58
+ )
59
+ return app
60
+
61
+
62
+class MockModelConfig:
63
+ def __init__(self):
64
+ self.hf_config = Namespace()
65
+ self.hf_config.num_labels = 1
66
+
67
+ def get_pooling_task(self, supported_tasks: tuple["SupportedTask", ...]):
68
+ pooling_tasks = [s for s in supported_tasks if s in POOLING_TASKS]
69
+ return pooling_tasks[0] if len(pooling_tasks) > 0 else None
70
+
71
+
72
+@pytest.fixture(params=get_args(SupportedTask))
73
+def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]]:
74
+ """For each supported task, build an app with only that task's routers,
75
+ extract all routes, and return the task name and routes."""
76
+ task = request.param
77
+ # Enable development mode to register all routes (including dev-only routes).
78
+ monkeypatch.setenv("VLLM_SERVER_DEV_MODE", "1")
79
+
80
+ app = FastAPI()
81
+ args = Namespace()
82
+ app.state = Namespace()
83
+ app.state.args = args
84
+
85
+ # Register routers for this specific task (development mode already enabled).
86
+ register_api_routers(
87
+ args, app, supported_tasks=(task,), model_config=MockModelConfig()
88
+ )
89
+
90
+ routes = get_all_http_routes(app)
91
+ return task, routes
92
+
93
+
94
+# ---------------------------------------------------------------------------
95
+# Tests for auto-discovered routes
96
+# ---------------------------------------------------------------------------
97
+
98
+
99
+def test_auto_discovered_protected_routes_require_auth(task_routes):
100
+ """For every auto-discovered route that starts with a guarded prefix,
101
+ verify that authentication is enforced."""
102
+ task, routes = task_routes
103
+ app = _create_app_with_mock_routes(routes)
104
+ client = TestClient(app)
105
+
106
+ for path_template, methods in routes:
107
+ if not path_template.startswith(GUARDED_PREFIX):
108
+ continue
109
+
110
+ test_path = generate_test_path(path_template)
111
+ test_method = methods[0] if methods else "GET"
112
+
113
+ resp = client.request(test_method, test_path)
114
+ assert resp.status_code == 401, (
115
+ f"[{task}] {test_method} {test_path} should reject missing token"
116
+ )
117
+
118
+ resp = client.request(
119
+ test_method, test_path, headers={"Authorization": "Bearer wrong"}
120
+ )
121
+ assert resp.status_code == 401, (
122
+ f"[{task}] {test_method} {test_path} should reject invalid token"
123
+ )
124
+
125
+ resp = client.request(
126
+ test_method, test_path, headers={"Authorization": "Bearer valid-token"}
127
+ )
128
+ assert resp.status_code == 200, (
129
+ f"[{task}] {test_method} {test_path} should accept valid token"
130
+ )
131
+
132
+
133
+def test_auto_discovered_unprotected_routes_no_auth(task_routes):
134
+ """For every auto-discovered route that does NOT start with a guarded
135
+ prefix, verify that no authentication is required."""
136
+ task, routes = task_routes
137
+ app = _create_app_with_mock_routes(routes)
138
+ client = TestClient(app)
139
+
140
+ for path_template, methods in routes:
141
+ if path_template.startswith(GUARDED_PREFIX):
142
+ continue
143
+
144
+ test_path = generate_test_path(path_template)
145
+ test_method = methods[0] if methods else "GET"
146
+
147
+ resp = client.request(test_method, test_path)
148
+ assert resp.status_code == 200, (
149
+ f"[{task}] {test_method} {test_path} should be accessible without token"
150
+ )
tests/entrypoints/serve/middleware/test_optional_middleware.py
renamed
vllm/entrypoints/launchers/__init__.py
vllm/entrypoints/launchers/api_server/__init__.py
vllm/entrypoints/launchers/api_server/routers.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from argparse import Namespace
4
+
5
+from fastapi import FastAPI
6
+
7
+from vllm import envs
8
+from vllm.config import ModelConfig
9
+from vllm.tasks import POOLING_TASKS, SupportedTask
10
+
11
+
12
+def register_api_routers(
13
+ args: Namespace,
14
+ app: FastAPI,
15
+ supported_tasks: tuple["SupportedTask", ...],
16
+ model_config: ModelConfig | None = None,
17
+):
18
+ from vllm.entrypoints.serve import register_vllm_serve_api_routers
19
+
20
+ register_vllm_serve_api_routers(app)
21
+
22
+ from vllm.entrypoints.openai.models.api_router import (
23
+ attach_router as register_models_api_router,
24
+ )
25
+
26
+ register_models_api_router(app)
27
+
28
+ from vllm.entrypoints.serve.sagemaker.api_router import (
29
+ attach_router as register_sagemaker_api_router,
30
+ )
31
+
32
+ register_sagemaker_api_router(app, supported_tasks, model_config)
33
+
34
+ if envs.VLLM_SERVER_DEV_MODE:
35
+ from vllm.entrypoints.serve import register_vllm_dev_api_routers
36
+
37
+ register_vllm_dev_api_routers(app)
38
+
39
+ if "generate" in supported_tasks:
40
+ from vllm.entrypoints.generate.api_router import (
41
+ register_generate_api_routers,
42
+ )
43
+
44
+ register_generate_api_routers(app)
45
+
46
+ from vllm.entrypoints.serve.elastic_ep.api_router import (
47
+ attach_router as elastic_ep_attach_router,
48
+ )
49
+
50
+ elastic_ep_attach_router(app)
51
+
52
+ if "generate" in supported_tasks or "render" in supported_tasks:
53
+ from vllm.entrypoints.scale_out.factories import register_scale_out_api_routers
54
+
55
+ register_scale_out_api_routers(app, supported_tasks)
56
+
57
+ if "transcription" in supported_tasks or "realtime" in supported_tasks:
58
+ from vllm.entrypoints.speech_to_text.factories import (
59
+ register_speech_to_text_api_routers,
60
+ )
61
+
62
+ register_speech_to_text_api_routers(app, supported_tasks)
63
+
64
+ if any(task in POOLING_TASKS for task in supported_tasks):
65
+ from vllm.entrypoints.pooling.factories import register_pooling_api_routers
66
+
67
+ register_pooling_api_routers(app, supported_tasks, model_config)
68
+
69
+ if getattr(args, "enable_fault_tolerance", False):
70
+ from vllm.entrypoints.serve.fault_tolerance.api_router import (
71
+ register_fault_tolerance_api_router,
72
+ )
73
+
74
+ register_fault_tolerance_api_router(app)
vllm/entrypoints/openai/api_server.py
+5
-114
@@ -1,8 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import asyncio
4
-import importlib
5
-import inspect
4
import multiprocessing
5
import multiprocessing.forkserver as forkserver
6
import os
@@ -17,7 +15,6 @@ from typing import Any, cast
15
16
import uvloop
17
from fastapi import FastAPI
20
-from fastapi.middleware.cors import CORSMiddleware
18
from starlette.datastructures import State
19
20
import vllm.envs as envs
@@ -26,11 +23,12 @@ from vllm.engine.arg_utils import AsyncEngineArgs
23
from vllm.engine.protocol import EngineClient
24
from vllm.entrypoints.chat_utils import load_chat_template
25
from vllm.entrypoints.launcher import serve_http
26
+from vllm.entrypoints.launchers.api_server.routers import register_api_routers
27
from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
28
from vllm.entrypoints.openai.models.protocol import BaseModelPath
29
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
32
-from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware
30
from vllm.entrypoints.serve.exception_handling.register import init_exception_handler
31
+from vllm.entrypoints.serve.middleware.register import init_entrypoints_middleware
32
from vllm.entrypoints.serve.sagemaker.api_router import sagemaker_standards_bootstrap
33
from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
34
from vllm.entrypoints.serve.utils.api_utils import (
@@ -43,7 +41,6 @@ from vllm.entrypoints.serve.utils.request_logger import RequestLogger
41
from vllm.entrypoints.serve.utils.server_utils import (
42
get_uvicorn_log_config,
43
lifespan,
46
- log_response,
44
)
45
from vllm.logger import init_logger
46
from vllm.reasoning import ReasoningParserManager
@@ -205,64 +202,9 @@ def build_app(
202
else:
203
app = FastAPI(lifespan=lifespan)
204
app.state.args = args
205
+ app.root_path = args.root_path
206
209
- from vllm.entrypoints.serve import register_vllm_serve_api_routers
210
-
211
- register_vllm_serve_api_routers(app)
212
-
213
- from vllm.entrypoints.openai.models.api_router import (
214
- attach_router as register_models_api_router,
215
- )
216
-
217
- register_models_api_router(app)
218
-
219
- from vllm.entrypoints.serve.sagemaker.api_router import (
220
- attach_router as register_sagemaker_api_router,
221
- )
222
-
223
- register_sagemaker_api_router(app, supported_tasks, model_config)
224
-
225
- if envs.VLLM_SERVER_DEV_MODE:
226
- from vllm.entrypoints.serve import register_vllm_dev_api_routers
227
-
228
- register_vllm_dev_api_routers(app)
229
-
230
- if "generate" in supported_tasks:
231
- from vllm.entrypoints.generate.api_router import (
232
- register_generate_api_routers,
233
- )
234
-
235
- register_generate_api_routers(app)
236
-
237
- from vllm.entrypoints.serve.elastic_ep.api_router import (
238
- attach_router as elastic_ep_attach_router,
239
- )
240
-
241
- elastic_ep_attach_router(app)
242
-
243
- if "generate" in supported_tasks or "render" in supported_tasks:
244
- from vllm.entrypoints.scale_out.factories import register_scale_out_api_routers
245
-
246
- register_scale_out_api_routers(app, supported_tasks)
247
-
248
- if "transcription" in supported_tasks or "realtime" in supported_tasks:
249
- from vllm.entrypoints.speech_to_text.factories import (
250
- register_speech_to_text_api_routers,
251
- )
252
-
253
- register_speech_to_text_api_routers(app, supported_tasks)
254
-
255
- if any(task in POOLING_TASKS for task in supported_tasks):
256
- from vllm.entrypoints.pooling.factories import register_pooling_api_routers
257
-
258
- register_pooling_api_routers(app, supported_tasks, model_config)
259
-
260
- if args.enable_fault_tolerance:
261
- from vllm.entrypoints.serve.fault_tolerance.api_router import (
262
- register_fault_tolerance_api_router,
263
- )
264
-
265
- register_fault_tolerance_api_router(app)
207
+ register_api_routers(args, app, supported_tasks, model_config)
208
209
# Endpoint plugins are attached last so their routes are registered after all core
210
# routers. This runs even for the CPU only render server. A plugin eligible for
@@ -271,58 +213,7 @@ def build_app(
213
_attach_endpoint_plugins(app, supported_tasks)
214
215
init_exception_handler(app)
274
-
275
- app.root_path = args.root_path
276
- app.add_middleware(
277
- CORSMiddleware,
278
- allow_origins=args.allowed_origins,
279
- allow_credentials=args.allow_credentials,
280
- allow_methods=args.allowed_methods,
281
- allow_headers=args.allowed_headers,
282
- )
283
-
284
- # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY
285
- if tokens := [key for key in (args.api_key or [envs.VLLM_API_KEY]) if key]:
286
- from vllm.entrypoints.serve.utils.server_utils import AuthenticationMiddleware
287
-
288
- app.add_middleware(AuthenticationMiddleware, tokens=tokens)
289
-
290
- if args.enable_request_id_headers:
291
- from vllm.entrypoints.serve.utils.server_utils import XRequestIdMiddleware
292
-
293
- app.add_middleware(XRequestIdMiddleware)
294
-
295
- # Add scaling middleware to check for scaling state
296
- app.add_middleware(ScalingMiddleware)
297
-
298
- if "realtime" in supported_tasks:
299
- # Add WebSocket metrics middleware
300
- from vllm.entrypoints.speech_to_text.factories import (
301
- add_websocket_metrics_middleware,
302
- )
303
-
304
- add_websocket_metrics_middleware(app)
305
-
306
- if envs.VLLM_DEBUG_LOG_API_SERVER_RESPONSE:
307
- logger.warning(
308
- "CAUTION: Enabling log response in the API Server. "
309
- "This can include sensitive information and should be "
310
- "avoided in production."
311
- )
312
- app.middleware("http")(log_response)
313
-
314
- for middleware in args.middleware:
315
- module_path, object_name = middleware.rsplit(".", 1)
316
- imported = getattr(importlib.import_module(module_path), object_name)
317
- if inspect.isclass(imported):
318
- app.add_middleware(imported) # type: ignore[arg-type]
319
- elif inspect.iscoroutinefunction(imported):
320
- app.middleware("http")(imported)
321
- else:
322
- raise ValueError(
323
- f"Invalid middleware {middleware}. Must be a function or a class."
324
- )
325
-
216
+ init_entrypoints_middleware(args, app, supported_tasks)
217
app = sagemaker_standards_bootstrap(app)
218
return app
219
vllm/entrypoints/serve/middleware/__init__.py
vllm/entrypoints/serve/middleware/authenticate.py
new
+62
@@ -0,0 +1,62 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import hashlib
4
+import secrets
5
+from collections.abc import Awaitable
6
+
7
+from starlette.datastructures import Headers
8
+from starlette.responses import JSONResponse
9
+from starlette.types import ASGIApp, Receive, Scope, Send
10
+
11
+GUARDED_PREFIX = ("/v1", "/v2", "/inference", "/cohere")
12
+
13
+
14
+class AuthenticationMiddleware:
15
+ """
16
+ Pure ASGI middleware that authenticates each request by checking
17
+ if the Authorization Bearer token exists and equals anyof "{api_key}".
18
+
19
+ Notes
20
+ -----
21
+ There are two cases in which authentication is skipped:
22
+ 1. The HTTP method is OPTIONS.
23
+ 2. The request path doesn't start with GUARDED_PREFIX (e.g. /health).
24
+ """
25
+
26
+ def __init__(self, app: ASGIApp, tokens: list[str]) -> None:
27
+ self.app = app
28
+ self.api_tokens = [hashlib.sha256(t.encode("utf-8")).digest() for t in tokens]
29
+
30
+ def verify_token(self, headers: Headers) -> bool:
31
+ authorization_header_value = headers.get("Authorization")
32
+ if not authorization_header_value:
33
+ return False
34
+
35
+ scheme, _, param = authorization_header_value.partition(" ")
36
+ if scheme.lower() != "bearer":
37
+ return False
38
+
39
+ param_hash = hashlib.sha256(param.encode("utf-8")).digest()
40
+
41
+ token_match = False
42
+ for token_hash in self.api_tokens:
43
+ token_match |= secrets.compare_digest(param_hash, token_hash)
44
+
45
+ return token_match
46
+
47
+ def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]:
48
+ if (
49
+ scope["type"] not in ("http", "websocket")
50
+ or scope.get("method") == "OPTIONS"
51
+ ):
52
+ # scope["type"] can be "lifespan" or "startup" for example,
53
+ # in which case we don't need to do anything
54
+ return self.app(scope, receive, send)
55
+ root_path = scope.get("root_path", "")
56
+ url_path = scope["path"].removeprefix(root_path)
57
+ headers = Headers(scope=scope)
58
+ # Type narrow to satisfy mypy.
59
+ if url_path.startswith(GUARDED_PREFIX) and not self.verify_token(headers):
60
+ response = JSONResponse(content={"error": "Unauthorized"}, status_code=401)
61
+ return response(scope, receive, send)
62
+ return self.app(scope, receive, send)
vllm/entrypoints/serve/middleware/log_response.py
new
+162
@@ -0,0 +1,162 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pydantic
4
+from fastapi import Request
5
+from starlette.concurrency import iterate_in_threadpool
6
+
7
+from vllm.logger import init_logger
8
+
9
+logger = init_logger(__name__)
10
+
11
+
12
+class SSEDecoder:
13
+ """Robust Server-Sent Events decoder for streaming responses."""
14
+
15
+ def __init__(self):
16
+ self.buffer = ""
17
+ self.content_buffer = []
18
+
19
+ def decode_chunk(self, chunk: bytes) -> list[dict]:
20
+ """Decode a chunk of SSE data and return parsed events."""
21
+ import json
22
+
23
+ try:
24
+ chunk_str = chunk.decode("utf-8")
25
+ except UnicodeDecodeError:
26
+ # Skip malformed chunks
27
+ return []
28
+
29
+ self.buffer += chunk_str
30
+ events = []
31
+
32
+ # Process complete lines
33
+ while "\n" in self.buffer:
34
+ line, self.buffer = self.buffer.split("\n", 1)
35
+ line = line.rstrip("\r") # Handle CRLF
36
+
37
+ if line.startswith("data: "):
38
+ data_str = line[6:].strip()
39
+ if data_str == "[DONE]":
40
+ events.append({"type": "done"})
41
+ elif data_str:
42
+ try:
43
+ event_data = json.loads(data_str)
44
+ events.append({"type": "data", "data": event_data})
45
+ except json.JSONDecodeError:
46
+ # Skip malformed JSON
47
+ continue
48
+
49
+ return events
50
+
51
+ def extract_content(self, event_data: dict) -> str:
52
+ """Extract content from event data."""
53
+ return _extract_content_from_chunk(event_data)
54
+
55
+ def add_content(self, content: str) -> None:
56
+ """Add content to the buffer."""
57
+ if content:
58
+ self.content_buffer.append(content)
59
+
60
+ def get_complete_content(self) -> str:
61
+ """Get the complete buffered content."""
62
+ return "".join(self.content_buffer)
63
+
64
+
65
+def _extract_content_from_chunk(chunk_data: dict) -> str:
66
+ """Extract content from a streaming response chunk."""
67
+ try:
68
+ from vllm.entrypoints.openai.chat_completion.protocol import (
69
+ ChatCompletionStreamResponse,
70
+ )
71
+ from vllm.entrypoints.openai.completion.protocol import (
72
+ CompletionStreamResponse,
73
+ )
74
+
75
+ # Try using Completion types for type-safe parsing
76
+ if chunk_data.get("object") == "chat.completion.chunk":
77
+ chat_response = ChatCompletionStreamResponse.model_validate(chunk_data)
78
+ if chat_response.choices and chat_response.choices[0].delta.content:
79
+ return chat_response.choices[0].delta.content
80
+ elif chunk_data.get("object") == "text_completion":
81
+ completion_response = CompletionStreamResponse.model_validate(chunk_data)
82
+ if completion_response.choices and completion_response.choices[0].text:
83
+ return completion_response.choices[0].text
84
+ except pydantic.ValidationError:
85
+ # Fallback to manual parsing
86
+ if "choices" in chunk_data and chunk_data["choices"]:
87
+ choice = chunk_data["choices"][0]
88
+ if "delta" in choice and choice["delta"].get("content"):
89
+ return choice["delta"]["content"]
90
+ elif choice.get("text"):
91
+ return choice["text"]
92
+ return ""
93
+
94
+
95
+def _log_streaming_response(response, response_body: list) -> None:
96
+ """Log streaming response with robust SSE parsing."""
97
+
98
+ sse_decoder = SSEDecoder()
99
+ chunk_count = 0
100
+
101
+ def buffered_iterator():
102
+ nonlocal chunk_count
103
+
104
+ for chunk in response_body:
105
+ chunk_count += 1
106
+ yield chunk
107
+
108
+ # Parse SSE events from chunk
109
+ events = sse_decoder.decode_chunk(chunk)
110
+
111
+ for event in events:
112
+ if event["type"] == "data":
113
+ content = sse_decoder.extract_content(event["data"])
114
+ sse_decoder.add_content(content)
115
+ elif event["type"] == "done":
116
+ # Log complete content when done
117
+ full_content = sse_decoder.get_complete_content()
118
+ if full_content:
119
+ # Truncate if too long
120
+ if len(full_content) > 2048:
121
+ full_content = full_content[:2048] + "...[truncated]"
122
+ logger.info(
123
+ "response_body={streaming_complete: content=%r, chunks=%d}",
124
+ full_content,
125
+ chunk_count,
126
+ )
127
+ else:
128
+ logger.info(
129
+ "response_body={streaming_complete: no_content, chunks=%d}",
130
+ chunk_count,
131
+ )
132
+ return
133
+
134
+ response.body_iterator = iterate_in_threadpool(buffered_iterator())
135
+ logger.info("response_body={streaming_started: chunks=%d}", len(response_body))
136
+
137
+
138
+def _log_non_streaming_response(response_body: list) -> None:
139
+ """Log non-streaming response."""
140
+ try:
141
+ decoded_body = response_body[0].decode()
142
+ logger.info("response_body={%s}", decoded_body)
143
+ except UnicodeDecodeError:
144
+ logger.info("response_body={<binary_data>}")
145
+
146
+
147
+async def log_response(request: Request, call_next):
148
+ response = await call_next(request)
149
+ response_body = [section async for section in response.body_iterator]
150
+ response.body_iterator = iterate_in_threadpool(iter(response_body))
151
+ # Check if this is a streaming response by looking at content-type
152
+ content_type = response.headers.get("content-type", "")
153
+ is_streaming = content_type == "text/event-stream; charset=utf-8"
154
+
155
+ # Log response body based on type
156
+ if not response_body:
157
+ logger.info("response_body={<empty>}")
158
+ elif is_streaming:
159
+ _log_streaming_response(response, response_body)
160
+ else:
161
+ _log_non_streaming_response(response_body)
162
+ return response
vllm/entrypoints/serve/middleware/register.py
new
+75
@@ -0,0 +1,75 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import importlib
4
+import inspect
5
+from argparse import Namespace
6
+
7
+from fastapi import FastAPI
8
+from starlette.middleware.cors import CORSMiddleware
9
+
10
+from vllm import envs
11
+from vllm.logger import init_logger
12
+from vllm.tasks import SupportedTask
13
+
14
+from .log_response import log_response
15
+
16
+logger = init_logger(__name__)
17
+
18
+
19
+def init_entrypoints_middleware(
20
+ args: Namespace,
21
+ app: FastAPI,
22
+ supported_tasks: tuple["SupportedTask", ...],
23
+):
24
+ app.add_middleware(
25
+ CORSMiddleware,
26
+ allow_origins=args.allowed_origins,
27
+ allow_credentials=args.allow_credentials,
28
+ allow_methods=args.allowed_methods,
29
+ allow_headers=args.allowed_headers,
30
+ )
31
+
32
+ # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY
33
+ if tokens := [key for key in (args.api_key or [envs.VLLM_API_KEY]) if key]:
34
+ from .authenticate import AuthenticationMiddleware
35
+
36
+ app.add_middleware(AuthenticationMiddleware, tokens=tokens)
37
+
38
+ if args.enable_request_id_headers:
39
+ from .x_request_id import XRequestIdMiddleware
40
+
41
+ app.add_middleware(XRequestIdMiddleware)
42
+
43
+ if "generate" in supported_tasks:
44
+ # Add scaling middleware to check for scaling state
45
+ from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware
46
+
47
+ app.add_middleware(ScalingMiddleware)
48
+
49
+ if "realtime" in supported_tasks:
50
+ # Add WebSocket metrics middleware
51
+ from vllm.entrypoints.speech_to_text.realtime.metrics import (
52
+ WebSocketMetricsMiddleware,
53
+ )
54
+
55
+ app.add_middleware(WebSocketMetricsMiddleware)
56
+
57
+ if envs.VLLM_DEBUG_LOG_API_SERVER_RESPONSE:
58
+ logger.warning(
59
+ "CAUTION: Enabling log response in the API Server. "
60
+ "This can include sensitive information and should be "
61
+ "avoided in production."
62
+ )
63
+ app.middleware("http")(log_response)
64
+
65
+ for middleware in args.middleware:
66
+ module_path, object_name = middleware.rsplit(".", 1)
67
+ imported = getattr(importlib.import_module(module_path), object_name)
68
+ if inspect.isclass(imported):
69
+ app.add_middleware(imported) # type: ignore[arg-type]
70
+ elif inspect.iscoroutinefunction(imported):
71
+ app.middleware("http")(imported)
72
+ else:
73
+ raise ValueError(
74
+ f"Invalid middleware {middleware}. Must be a function or a class."
75
+ )
vllm/entrypoints/serve/middleware/x_request_id.py
new
+38
@@ -0,0 +1,38 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import uuid
4
+from collections.abc import Awaitable
5
+
6
+from starlette.datastructures import Headers, MutableHeaders
7
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
8
+
9
+
10
+class XRequestIdMiddleware:
11
+ """
12
+ Middleware the set's the X-Request-Id header for each response
13
+ to a random uuid4 (hex) value if the header isn't already
14
+ present in the request, otherwise use the provided request id.
15
+ """
16
+
17
+ def __init__(self, app: ASGIApp) -> None:
18
+ self.app = app
19
+
20
+ def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]:
21
+ if scope["type"] not in ("http", "websocket"):
22
+ return self.app(scope, receive, send)
23
+
24
+ # Extract the request headers.
25
+ request_headers = Headers(scope=scope)
26
+
27
+ async def send_with_request_id(message: Message) -> None:
28
+ """
29
+ Custom send function to mutate the response headers
30
+ and append X-Request-Id to it.
31
+ """
32
+ if message["type"] == "http.response.start":
33
+ response_headers = MutableHeaders(raw=message["headers"])
34
+ request_id = request_headers.get("X-Request-Id", uuid.uuid4().hex)
35
+ response_headers.append("X-Request-Id", request_id)
36
+ await send(message)
37
+
38
+ return self.app(scope, receive, send_with_request_id)
vllm/entrypoints/serve/utils/server_utils.py
+1
-250
@@ -1,20 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import asyncio
4
-import hashlib
4
import json
6
-import secrets
7
-import uuid
5
from argparse import Namespace
9
-from collections.abc import Awaitable
6
from contextlib import asynccontextmanager
7
12
-import pydantic
13
-from fastapi import FastAPI, Request
14
-from fastapi.responses import JSONResponse
15
-from starlette.concurrency import iterate_in_threadpool
16
-from starlette.datastructures import Headers, MutableHeaders
17
-from starlette.types import ASGIApp, Message, Receive, Scope, Send
8
+from fastapi import FastAPI
9
10
from vllm import envs
11
from vllm.engine.protocol import EngineClient
@@ -24,91 +15,6 @@ from vllm.utils.gc_utils import freeze_gc_heap
15
logger = init_logger("vllm.entrypoints.openai.server_utils")
16
17
27
-GUARDED_PREFIX = ("/v1", "/v2", "/inference", "/cohere")
28
-
29
-
30
-class AuthenticationMiddleware:
31
- """
32
- Pure ASGI middleware that authenticates each request by checking
33
- if the Authorization Bearer token exists and equals anyof "{api_key}".
34
-
35
- Notes
36
- -----
37
- There are two cases in which authentication is skipped:
38
- 1. The HTTP method is OPTIONS.
39
- 2. The request path doesn't start with GUARDED_PREFIX (e.g. /health).
40
- """
41
-
42
- def __init__(self, app: ASGIApp, tokens: list[str]) -> None:
43
- self.app = app
44
- self.api_tokens = [hashlib.sha256(t.encode("utf-8")).digest() for t in tokens]
45
-
46
- def verify_token(self, headers: Headers) -> bool:
47
- authorization_header_value = headers.get("Authorization")
48
- if not authorization_header_value:
49
- return False
50
-
51
- scheme, _, param = authorization_header_value.partition(" ")
52
- if scheme.lower() != "bearer":
53
- return False
54
-
55
- param_hash = hashlib.sha256(param.encode("utf-8")).digest()
56
-
57
- token_match = False
58
- for token_hash in self.api_tokens:
59
- token_match |= secrets.compare_digest(param_hash, token_hash)
60
-
61
- return token_match
62
-
63
- def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]:
64
- if (
65
- scope["type"] not in ("http", "websocket")
66
- or scope.get("method") == "OPTIONS"
67
- ):
68
- # scope["type"] can be "lifespan" or "startup" for example,
69
- # in which case we don't need to do anything
70
- return self.app(scope, receive, send)
71
- root_path = scope.get("root_path", "")
72
- url_path = scope["path"].removeprefix(root_path)
73
- headers = Headers(scope=scope)
74
- # Type narrow to satisfy mypy.
75
- if url_path.startswith(GUARDED_PREFIX) and not self.verify_token(headers):
76
- response = JSONResponse(content={"error": "Unauthorized"}, status_code=401)
77
- return response(scope, receive, send)
78
- return self.app(scope, receive, send)
79
-
80
-
81
-class XRequestIdMiddleware:
82
- """
83
- Middleware the set's the X-Request-Id header for each response
84
- to a random uuid4 (hex) value if the header isn't already
85
- present in the request, otherwise use the provided request id.
86
- """
87
-
88
- def __init__(self, app: ASGIApp) -> None:
89
- self.app = app
90
-
91
- def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]:
92
- if scope["type"] not in ("http", "websocket"):
93
- return self.app(scope, receive, send)
94
-
95
- # Extract the request headers.
96
- request_headers = Headers(scope=scope)
97
-
98
- async def send_with_request_id(message: Message) -> None:
99
- """
100
- Custom send function to mutate the response headers
101
- and append X-Request-Id to it.
102
- """
103
- if message["type"] == "http.response.start":
104
- response_headers = MutableHeaders(raw=message["headers"])
105
- request_id = request_headers.get("X-Request-Id", uuid.uuid4().hex)
106
- response_headers.append("X-Request-Id", request_id)
107
- await send(message)
108
-
109
- return self.app(scope, receive, send_with_request_id)
110
-
111
-
18
def load_log_config(log_config_file: str | None) -> dict | None:
19
if not log_config_file:
20
return None
@@ -155,161 +61,6 @@ def get_uvicorn_log_config(args: Namespace) -> dict | None:
61
return None
62
63
158
-def _extract_content_from_chunk(chunk_data: dict) -> str:
159
- """Extract content from a streaming response chunk."""
160
- try:
161
- from vllm.entrypoints.openai.chat_completion.protocol import (
162
- ChatCompletionStreamResponse,
163
- )
164
- from vllm.entrypoints.openai.completion.protocol import (
165
- CompletionStreamResponse,
166
- )
167
-
168
- # Try using Completion types for type-safe parsing
169
- if chunk_data.get("object") == "chat.completion.chunk":
170
- chat_response = ChatCompletionStreamResponse.model_validate(chunk_data)
171
- if chat_response.choices and chat_response.choices[0].delta.content:
172
- return chat_response.choices[0].delta.content
173
- elif chunk_data.get("object") == "text_completion":
174
- completion_response = CompletionStreamResponse.model_validate(chunk_data)
175
- if completion_response.choices and completion_response.choices[0].text:
176
- return completion_response.choices[0].text
177
- except pydantic.ValidationError:
178
- # Fallback to manual parsing
179
- if "choices" in chunk_data and chunk_data["choices"]:
180
- choice = chunk_data["choices"][0]
181
- if "delta" in choice and choice["delta"].get("content"):
182
- return choice["delta"]["content"]
183
- elif choice.get("text"):
184
- return choice["text"]
185
- return ""
186
-
187
-
188
-class SSEDecoder:
189
- """Robust Server-Sent Events decoder for streaming responses."""
190
-
191
- def __init__(self):
192
- self.buffer = ""
193
- self.content_buffer = []
194
-
195
- def decode_chunk(self, chunk: bytes) -> list[dict]:
196
- """Decode a chunk of SSE data and return parsed events."""
197
- import json
198
-
199
- try:
200
- chunk_str = chunk.decode("utf-8")
201
- except UnicodeDecodeError:
202
- # Skip malformed chunks
203
- return []
204
-
205
- self.buffer += chunk_str
206
- events = []
207
-
208
- # Process complete lines
209
- while "\n" in self.buffer:
210
- line, self.buffer = self.buffer.split("\n", 1)
211
- line = line.rstrip("\r") # Handle CRLF
212
-
213
- if line.startswith("data: "):
214
- data_str = line[6:].strip()
215
- if data_str == "[DONE]":
216
- events.append({"type": "done"})
217
- elif data_str:
218
- try:
219
- event_data = json.loads(data_str)
220
- events.append({"type": "data", "data": event_data})
221
- except json.JSONDecodeError:
222
- # Skip malformed JSON
223
- continue
224
-
225
- return events
226
-
227
- def extract_content(self, event_data: dict) -> str:
228
- """Extract content from event data."""
229
- return _extract_content_from_chunk(event_data)
230
-
231
- def add_content(self, content: str) -> None:
232
- """Add content to the buffer."""
233
- if content:
234
- self.content_buffer.append(content)
235
-
236
- def get_complete_content(self) -> str:
237
- """Get the complete buffered content."""
238
- return "".join(self.content_buffer)
239
-
240
-
241
-def _log_streaming_response(response, response_body: list) -> None:
242
- """Log streaming response with robust SSE parsing."""
243
- from starlette.concurrency import iterate_in_threadpool
244
-
245
- sse_decoder = SSEDecoder()
246
- chunk_count = 0
247
-
248
- def buffered_iterator():
249
- nonlocal chunk_count
250
-
251
- for chunk in response_body:
252
- chunk_count += 1
253
- yield chunk
254
-
255
- # Parse SSE events from chunk
256
- events = sse_decoder.decode_chunk(chunk)
257
-
258
- for event in events:
259
- if event["type"] == "data":
260
- content = sse_decoder.extract_content(event["data"])
261
- sse_decoder.add_content(content)
262
- elif event["type"] == "done":
263
- # Log complete content when done
264
- full_content = sse_decoder.get_complete_content()
265
- if full_content:
266
- # Truncate if too long
267
- if len(full_content) > 2048:
268
- full_content = full_content[:2048] + ""
269
- "...[truncated]"
270
- logger.info(
271
- "response_body={streaming_complete: content=%r, chunks=%d}",
272
- full_content,
273
- chunk_count,
274
- )
275
- else:
276
- logger.info(
277
- "response_body={streaming_complete: no_content, chunks=%d}",
278
- chunk_count,
279
- )
280
- return
281
-
282
- response.body_iterator = iterate_in_threadpool(buffered_iterator())
283
- logger.info("response_body={streaming_started: chunks=%d}", len(response_body))
284
-
285
-
286
-def _log_non_streaming_response(response_body: list) -> None:
287
- """Log non-streaming response."""
288
- try:
289
- decoded_body = response_body[0].decode()
290
- logger.info("response_body={%s}", decoded_body)
291
- except UnicodeDecodeError:
292
- logger.info("response_body={<binary_data>}")
293
-
294
-
295
-async def log_response(request: Request, call_next):
296
- response = await call_next(request)
297
- response_body = [section async for section in response.body_iterator]
298
- response.body_iterator = iterate_in_threadpool(iter(response_body))
299
- # Check if this is a streaming response by looking at content-type
300
- content_type = response.headers.get("content-type", "")
301
- is_streaming = content_type == "text/event-stream; charset=utf-8"
302
-
303
- # Log response body based on type
304
- if not response_body:
305
- logger.info("response_body={<empty>}")
306
- elif is_streaming:
307
- _log_streaming_response(response, response_body)
308
- else:
309
- _log_non_streaming_response(response_body)
310
- return response
311
-
312
-
64
_running_tasks: set[asyncio.Task] = set()
65
66
vllm/entrypoints/speech_to_text/factories.py
-6
@@ -37,12 +37,6 @@ def register_speech_to_text_api_routers(
37
app.include_router(translation_router)
38
39
40
-def add_websocket_metrics_middleware(app: FastAPI):
41
- from .realtime.metrics import WebSocketMetricsMiddleware
42
-
43
- app.add_middleware(WebSocketMetricsMiddleware)
44
-
45
-
40
def init_speech_to_text_state(
41
engine_client: "EngineClient",
42
state: "State",