본문으로 건너뛰기

[Open WebUI] APIKeyRestrictionMiddleware를 순수 ASGI로 전환하여 스트리밍 오버헤드 제거

PR 링크: open-webui/open-webui#22188 상태: Merged | 변경: +44 / -38

들어가며

Starlette의 BaseHTTPMiddleware는 편리하지만, 내부적으로 send 콜러블을 래핑합니다. 이는 스트리밍 응답에서 모든 아웃바운드 청크가 미들웨어의 Python 호출 스택을 통과한다는 것을 의미합니다. APIKeyRestrictionMiddleware는 요청을 거부(403)하거나 그대로 통과시키는 단순한 역할인데도, 통과 경로에서 불필요한 오버헤드가 발생하고 있었습니다.

핵심 코드 분석

기존: BaseHTTPMiddleware 상속

Before:

class APIKeyRestrictionMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        auth_header = request.headers.get("Authorization")
        token = None
        if auth_header:
            parts = auth_header.split(" ", 1)
            if len(parts) == 2:
                token = parts[1]

        if token and token.startswith("sk-"):
            # 제한 로직...
            if not is_allowed:
                return JSONResponse(status_code=403, content={...})

        response = await call_next(request)
        return response

변경: 순수 ASGI 미들웨어

After:

class APIKeyRestrictionMiddleware:
    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] == "http":
            request = Request(scope)
            auth_header = request.headers.get("Authorization")
            token = None
            if auth_header:
                parts = auth_header.split(" ", 1)
                if len(parts) == 2:
                    token = parts[1]

            if token and token.startswith("sk-"):
                # 제한 로직...
                if not is_allowed:
                    await JSONResponse(status_code=403, content={...})(scope, receive, send)
                    return

        await self.app(scope, receive, send)

왜 이게 좋은가

  • 스트리밍 오버헤드 제거: 통과 경로에서 send를 래핑하지 않으므로, 청크가 소켓으로 직접 전달됩니다.
  • 청크당 Python 호출 스택 제거: LLM 스트리밍 응답처럼 청크가 많은 경우 누적 효과가 큽니다.
  • 동일한 동작 보장: 경로 매칭, 응답 형식 등 기존 로직이 그대로 유지됩니다.
  • ASGI 표준 준수: BaseHTTPMiddleware의 알려진 한계를 우회하는 Starlette 권장 패턴입니다.

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글