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