From ee28032fb9bd830faca716a089e40429f3776eb8 Mon Sep 17 00:00:00 2001 From: Classic298 <27028174+Classic298@users.noreply.github.com> Date: Tue, 14 Apr 2026 17:47:48 +0200 Subject: [PATCH] fix(middleware): replace BaseHTTPMiddleware HTTP middlewares with pure ASGI implementations (#23709) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(middleware): replace BaseHTTPMiddleware HTTP middlewares with pure ASGI implementations Starlette's BaseHTTPMiddleware (and the @app.middleware('http') decorator that uses it) wraps the downstream app in an anyio task group whose cancel scope tears down the inner task on every exit — client disconnect, response complete, or any outer middleware bailing. That CancelledError gets injected into whatever the inner task was awaiting, so DB queries, embedding calls, and other long awaits get killed mid-flight. Under aiosqlite the cleanup path then logs a multi-page `terminate_force_close() not implemented` traceback at ERROR for every cancelled DB call. Open WebUI had four such middlewares stacked (`commit_session_after_request`, `check_url`, `inspect_websocket`, `RedirectMiddleware`) so a single cancellation would compound through all four. Move the four middlewares to a new `open_webui.utils.asgi_middleware` module as plain ASGI classes (`__call__(scope, receive, send)`): * `CommitSessionMiddleware` — was `commit_session_after_request`; now also rolls back if commit fails before releasing the connection. * `AuthTokenMiddleware` — was `check_url`; sets request.state token + enable_api_keys + stamps X-Process-Time via a wrapped send. * `WebsocketUpgradeGuardMiddleware` — was `inspect_websocket`; rejects /ws/socket.io HTTP requests that claim transport=websocket without a proper Upgrade/Connection header. * `RedirectMiddleware` — was the BaseHTTPMiddleware subclass; same /watch + share-target rewrites. Pure ASGI does not introduce a cancel scope around the downstream app, so client disconnects propagate via `receive()` (the way ASGI was designed) instead of being injected as CancelledError. Middleware ordering is preserved. https://claude.ai/code/session_01JSr4NZSskEUQvoJnavVXh8 * fix(middleware): CommitSessionMiddleware — rollback on downstream error, never commit failed requests The first cut put commit() in a finally block, which meant that even when a downstream handler raised, the middleware would still commit whatever partial sync writes that handler had made before the failure. That regressed the previous BaseHTTPMiddleware semantics where commit only ran on the success path. Restructure the failure handling: * Downstream raised → rollback any pending sync work, release the connection, re-raise so the outer error middleware turns it into an error response. We never commit a request that did not complete. * Downstream returned → commit. On commit failure, log loudly, rollback, and re-raise. ScopedSession.remove() always runs in finally so the connection cannot leak. Document the inherent pure-ASGI limitation explicitly: by the time `await self.app(...)` returns the response messages have already been emitted, so a commit failure can no longer change what the client sees on the wire. Buffering the response to gate it on commit success would break streaming responses (chat completions, SSE) which are core to Open WebUI; the trade-off is intentional. Routes that need commit-before-send must manage the sync session explicitly. Also drop unused `typing` imports flagged by review. https://claude.ai/code/session_01JSr4NZSskEUQvoJnavVXh8 --------- Co-authored-by: Claude --- backend/open_webui/main.py | 113 ++------ backend/open_webui/utils/asgi_middleware.py | 273 ++++++++++++++++++++ 2 files changed, 290 insertions(+), 96 deletions(-) create mode 100644 backend/open_webui/utils/asgi_middleware.py diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 71372d0495..500d874e0c 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -46,7 +46,6 @@ from fastapi.staticfiles import StaticFiles from starlette_compress import CompressMiddleware from starlette.exceptions import HTTPException as StarletteHTTPException -from starlette.middleware.base import BaseHTTPMiddleware from starlette.middleware.sessions import SessionMiddleware from starlette.responses import Response, StreamingResponse from starlette.datastructures import Headers @@ -58,6 +57,12 @@ from starsessions import ( from starsessions.stores.redis import RedisStore from open_webui.utils import logger +from open_webui.utils.asgi_middleware import ( + AuthTokenMiddleware, + CommitSessionMiddleware, + RedirectMiddleware, + WebsocketUpgradeGuardMiddleware, +) from open_webui.utils.audit import AuditLevel, AuditLoggingMiddleware from open_webui.utils.logger import start_logger from open_webui.utils.session_pool import get_session @@ -1359,103 +1364,19 @@ if ENABLE_COMPRESSION_MIDDLEWARE: app.add_middleware(CompressMiddleware) -class RedirectMiddleware(BaseHTTPMiddleware): - async def dispatch(self, request: Request, call_next): - # Check if the request is a GET request - if request.method == 'GET': - path = request.url.path - query_params = dict(parse_qs(urlparse(str(request.url)).query)) - - redirect_params = {} - - # Check for the specific watch path and the presence of 'v' parameter - if path.endswith('/watch') and 'v' in query_params: - # Extract the first 'v' parameter - youtube_video_id = query_params['v'][0] - redirect_params['youtube'] = youtube_video_id - - if 'shared' in query_params and len(query_params['shared']) > 0: - # PWA share_target support - - text = query_params['shared'][0] - if text: - urls = re.match(r'https://\S+', text) - if urls: - from open_webui.retrieval.loaders.youtube import _parse_video_id - - if youtube_video_id := _parse_video_id(urls[0]): - redirect_params['youtube'] = youtube_video_id - else: - redirect_params['load-url'] = urls[0] - else: - redirect_params['q'] = text - - if redirect_params: - redirect_url = f'/?{urlencode(redirect_params)}' - return RedirectResponse(url=redirect_url) - - # Proceed with the normal flow of other requests - response = await call_next(request) - return response - - +# All HTTP middlewares below are pure-ASGI implementations. The previous +# `BaseHTTPMiddleware` / `@app.middleware('http')` versions wrapped the +# downstream app in an anyio task group whose cancel scope cancelled +# in-flight DB calls (and any other awaits) on client disconnect / +# response completion — which surfaced as noisy SQLAlchemy +# `terminate_force_close` tracebacks under aiosqlite and as random +# CancelledError storms across the request path. See +# `open_webui.utils.asgi_middleware` for the rationale. app.add_middleware(RedirectMiddleware) app.add_middleware(SecurityHeadersMiddleware) - - -@app.middleware('http') -async def commit_session_after_request(request: Request, call_next): - response = await call_next(request) - # log.debug("Commit session after request") - try: - ScopedSession.commit() - finally: - # CRITICAL: remove() returns the connection to the pool. - # Without this, connections remain "checked out" and accumulate - # as "idle in transaction" in PostgreSQL. - ScopedSession.remove() - return response - - -@app.middleware('http') -async def check_url(request: Request, call_next): - start_time = int(time.time()) - request.state.token = get_http_authorization_cred(request.headers.get('Authorization')) - # Fallback to cookie token for browser sessions - if request.state.token is None and request.cookies.get('token'): - from fastapi.security import HTTPAuthorizationCredentials - - request.state.token = HTTPAuthorizationCredentials(scheme='Bearer', credentials=request.cookies.get('token')) - - # Fallback to x-api-key header (Anthropic-compatible clients use this - # for ALL requests, including GET /v1/models, not just POST /v1/messages). - if request.state.token is None and request.headers.get('x-api-key'): - from fastapi.security import HTTPAuthorizationCredentials - - request.state.token = HTTPAuthorizationCredentials( - scheme='Bearer', credentials=request.headers.get('x-api-key') - ) - - request.state.enable_api_keys = app.state.config.ENABLE_API_KEYS - response = await call_next(request) - process_time = int(time.time()) - start_time - response.headers['X-Process-Time'] = str(process_time) - return response - - -@app.middleware('http') -async def inspect_websocket(request: Request, call_next): - if '/ws/socket.io' in request.url.path and request.query_params.get('transport') == 'websocket': - upgrade = (request.headers.get('Upgrade') or '').lower() - connection = (request.headers.get('Connection') or '').lower().split(',') - # Check that there's the correct headers for an upgrade, else reject the connection - # This is to work around this upstream issue: https://github.com/miguelgrinberg/python-engineio/issues/367 - if upgrade != 'websocket' or 'upgrade' not in connection: - return JSONResponse( - status_code=status.HTTP_400_BAD_REQUEST, - content={'detail': 'Invalid WebSocket upgrade request'}, - ) - return await call_next(request) +app.add_middleware(CommitSessionMiddleware) +app.add_middleware(AuthTokenMiddleware, fastapi_app=app) +app.add_middleware(WebsocketUpgradeGuardMiddleware) app.add_middleware( diff --git a/backend/open_webui/utils/asgi_middleware.py b/backend/open_webui/utils/asgi_middleware.py new file mode 100644 index 0000000000..4d29a79e66 --- /dev/null +++ b/backend/open_webui/utils/asgi_middleware.py @@ -0,0 +1,273 @@ +""" +Pure-ASGI replacements for the project's previous +`@app.middleware('http')` / `BaseHTTPMiddleware` middlewares. + +Why this matters +---------------- +Starlette's `BaseHTTPMiddleware` (which `@app.middleware('http')` is +sugar for) runs the downstream app inside an `anyio` task group. When +the wrapper exits — for any reason: response complete, client +disconnect, an outer middleware bailing out — the task group cancels +the inner task. That `CancelledError` then propagates into whatever +the inner task was doing, including in-flight DB queries, embedding +calls and disk I/O. + +In Open WebUI this surfaces as: + +* SQLAlchemy logging multi-page `NotImplementedError: + terminate_force_close()` tracebacks at ERROR every time a request is + cancelled mid-DB-call (the aiosqlite connector cleanup path). +* Spurious cancellations cascading through the four stacked + `@app.middleware('http')` wrappers. + +Pure ASGI middleware does not introduce a cancel scope around the +downstream app, so client disconnects propagate the way ASGI was +designed to (via `receive()` returning `http.disconnect`) instead of +being injected as `CancelledError` into arbitrary `await` points. + +Reference: https://www.starlette.io/middleware/#limitations +""" + +from __future__ import annotations + +import logging +import re +import time +from urllib.parse import parse_qs, urlencode + +from fastapi.responses import JSONResponse, RedirectResponse +from fastapi.security import HTTPAuthorizationCredentials +from starlette.datastructures import MutableHeaders +from starlette.requests import Request +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from open_webui.internal.db import ScopedSession +from open_webui.utils.auth import get_http_authorization_cred + +log = logging.getLogger(__name__) + + +class CommitSessionMiddleware: + """Commit and release the thread-local sync `ScopedSession` after each + HTTP request. + + Most requests now use the async session; the sync ScopedSession is + only touched by startup, healthchecks, and a handful of legacy + helpers (notably the pgvector / opengauss vector-DB clients). The + middleware exists so that PostgreSQL connections do not accumulate + as "idle in transaction" and so that any pending sync work made + inside the request is durably persisted. + + Failure semantics + ----------------- + * Downstream raised → roll back any pending sync work, release the + connection, and re-raise so the outer exception middleware can + turn it into an error response. We never commit work on a + request that did not complete successfully. + * Downstream returned → commit pending sync work; on commit + failure, log loudly, roll back, and re-raise. Note that in pure + ASGI the response messages have already been emitted by the + time `await self.app(...)` returns, so a commit failure cannot + retroactively change what the client sees on the wire — but + re-raising still surfaces the error in logs and to ASGI servers + that expose it. We deliberately do not buffer the response to + gate it on commit success, because that would defeat streaming + responses (chat completions, SSE) which are core to the app. + + For request paths where commit-before-send is required, manage the + sync session explicitly inside the handler instead of relying on + this middleware. + """ + + def __init__(self, app: ASGIApp) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope['type'] != 'http': + await self.app(scope, receive, send) + return + + try: + await self.app(scope, receive, send) + except BaseException: + # Downstream did not complete successfully. Roll back any + # pending sync writes, release the connection, and let the + # exception propagate. + try: + ScopedSession.rollback() + except Exception: + log.exception('CommitSessionMiddleware: rollback failed after downstream error') + finally: + ScopedSession.remove() + raise + + # Downstream completed. Commit pending sync work. + try: + ScopedSession.commit() + except Exception: + log.exception( + 'CommitSessionMiddleware: post-request commit failed; ' + 'response was already sent to client' + ) + try: + ScopedSession.rollback() + except Exception: + log.exception('CommitSessionMiddleware: rollback failed after commit failure') + raise + finally: + # CRITICAL: remove() returns the connection to the pool. + # Without this, connections remain "checked out" and + # accumulate as "idle in transaction" in PostgreSQL. + ScopedSession.remove() + + +class AuthTokenMiddleware: + """Extract the bearer/cookie/x-api-key credential and stash it on + `request.state.token`. + + Routes that depend on `get_verified_user` etc. read this state. + Also exposes `request.state.enable_api_keys` (snapshotted at request + entry from runtime config) and stamps an `X-Process-Time` response + header. + """ + + def __init__(self, app: ASGIApp, *, fastapi_app) -> None: + self.app = app + self._fastapi_app = fastapi_app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope['type'] != 'http': + await self.app(scope, receive, send) + return + + start_time = time.monotonic() + request = Request(scope) + + token = get_http_authorization_cred(request.headers.get('Authorization')) + if token is None: + cookie_token = request.cookies.get('token') + if cookie_token: + token = HTTPAuthorizationCredentials(scheme='Bearer', credentials=cookie_token) + if token is None: + api_key = request.headers.get('x-api-key') + if api_key: + token = HTTPAuthorizationCredentials(scheme='Bearer', credentials=api_key) + + request.state.token = token + request.state.enable_api_keys = self._fastapi_app.state.config.ENABLE_API_KEYS + + async def send_with_timing(message: Message) -> None: + if message['type'] == 'http.response.start': + process_time = int(time.monotonic() - start_time) + headers = MutableHeaders(scope=message) + headers['X-Process-Time'] = str(process_time) + await send(message) + + await self.app(scope, receive, send_with_timing) + + +class WebsocketUpgradeGuardMiddleware: + """Reject HTTP requests to `/ws/socket.io` that claim + `transport=websocket` but lack the proper `Upgrade`/`Connection` + headers. + + Works around https://github.com/miguelgrinberg/python-engineio/issues/367 + where engineio mishandles such requests. + """ + + def __init__(self, app: ASGIApp) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope['type'] != 'http': + await self.app(scope, receive, send) + return + + path = scope.get('path', '') + if '/ws/socket.io' in path: + query_string = scope.get('query_string', b'').decode('latin-1', errors='replace') + query_params = parse_qs(query_string) + if query_params.get('transport', [''])[0] == 'websocket': + headers = _scope_headers(scope) + upgrade = headers.get('upgrade', '').lower() + connection_tokens = [ + token.strip() for token in headers.get('connection', '').lower().split(',') + ] + if upgrade != 'websocket' or 'upgrade' not in connection_tokens: + response = JSONResponse( + status_code=400, + content={'detail': 'Invalid WebSocket upgrade request'}, + ) + await response(scope, receive, send) + return + + await self.app(scope, receive, send) + + +class RedirectMiddleware: + """Rewrites a couple of legacy entry-points to the SPA's own routes: + + * ``GET /watch?v=ID`` (YouTube) → ``/?youtube=ID`` + * ``GET /?shared=…`` (PWA share-target) → ``/?youtube=…`` / + ``/?load-url=…`` / ``/?q=…`` + """ + + def __init__(self, app: ASGIApp) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope['type'] != 'http' or scope.get('method', '').upper() != 'GET': + await self.app(scope, receive, send) + return + + path = scope.get('path', '') + query_string = scope.get('query_string', b'').decode('latin-1', errors='replace') + query_params = parse_qs(query_string) + + redirect_params: dict[str, str] = {} + if path.endswith('/watch') and 'v' in query_params and query_params['v']: + redirect_params['youtube'] = query_params['v'][0] + + if 'shared' in query_params and query_params['shared']: + text = query_params['shared'][0] + if text: + url_match = re.match(r'https://\S+', text) + if url_match: + # Local import: youtube loader pulls heavy deps and is + # only needed when a share-target actually contains a + # YouTube URL. + from open_webui.retrieval.loaders.youtube import _parse_video_id + + youtube_video_id = _parse_video_id(url_match[0]) + if youtube_video_id: + redirect_params['youtube'] = youtube_video_id + else: + redirect_params['load-url'] = url_match[0] + else: + redirect_params['q'] = text + + if redirect_params: + redirect_url = f'/?{urlencode(redirect_params)}' + response = RedirectResponse(url=redirect_url) + await response(scope, receive, send) + return + + await self.app(scope, receive, send) + + +def _scope_headers(scope: Scope) -> dict[str, str]: + """Return ASGI scope headers as a lower-cased str→str dict. + + ASGI delivers headers as a list of (bytes, bytes) pairs. For + convenience, fold duplicate keys with comma-joining (matching + HTTP/1.1 semantics). + """ + decoded: dict[str, str] = {} + for raw_key, raw_value in scope.get('headers', []): + key = raw_key.decode('latin-1').lower() + value = raw_value.decode('latin-1') + if key in decoded: + decoded[key] = f'{decoded[key]}, {value}' + else: + decoded[key] = value + return decoded