mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-26 03:25:43 +08:00
refactor: move cache-controls to /middleware
We can test the code without excessively mocking server dependencies
This commit is contained in:
parent
180e19e923
commit
a0e084042e
1
middleware/__init__.py
Normal file
1
middleware/__init__.py
Normal file
@ -0,0 +1 @@
|
|||||||
|
"""Server middleware modules"""
|
||||||
46
middleware/cache_middleware.py
Normal file
46
middleware/cache_middleware.py
Normal file
@ -0,0 +1,46 @@
|
|||||||
|
"""Cache control middleware for ComfyUI server"""
|
||||||
|
|
||||||
|
from aiohttp import web
|
||||||
|
from typing import Callable, Awaitable
|
||||||
|
|
||||||
|
# Time in seconds
|
||||||
|
ONE_HOUR: int = 3600
|
||||||
|
ONE_DAY: int = 86400
|
||||||
|
IMG_EXTENSIONS = (
|
||||||
|
".jpg",
|
||||||
|
".jpeg",
|
||||||
|
".png",
|
||||||
|
".ppm",
|
||||||
|
".bmp",
|
||||||
|
".pgm",
|
||||||
|
".tif",
|
||||||
|
".tiff",
|
||||||
|
".webp",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@web.middleware
|
||||||
|
async def cache_control(
|
||||||
|
request: web.Request, handler: Callable[[web.Request], Awaitable[web.Response]]
|
||||||
|
) -> web.Response:
|
||||||
|
"""Cache control middleware that sets appropriate cache headers based on file type and response status"""
|
||||||
|
response: web.Response = await handler(request)
|
||||||
|
|
||||||
|
if (
|
||||||
|
request.path.endswith(".js")
|
||||||
|
or request.path.endswith(".css")
|
||||||
|
or request.path.endswith("index.json")
|
||||||
|
):
|
||||||
|
response.headers.setdefault("Cache-Control", "no-cache")
|
||||||
|
elif request.path.lower().endswith(IMG_EXTENSIONS):
|
||||||
|
if response.status == 404:
|
||||||
|
response.headers.setdefault("Cache-Control", f"public, max-age={ONE_HOUR}")
|
||||||
|
elif response.status in (200, 201, 202, 203, 204, 205, 206, 301, 308):
|
||||||
|
# Success responses and permanent redirects - cache for 1 day
|
||||||
|
response.headers.setdefault("Cache-Control", f"public, max-age={ONE_DAY}")
|
||||||
|
elif response.status in (302, 303, 307):
|
||||||
|
# Temporary redirects - no cache
|
||||||
|
response.headers.setdefault("Cache-Control", "no-cache")
|
||||||
|
# Note: 304 Not Modified falls through - no cache headers set
|
||||||
|
|
||||||
|
return response
|
||||||
23
server.py
23
server.py
@ -39,10 +39,8 @@ from typing import Optional, Union, Callable, Awaitable
|
|||||||
from api_server.routes.internal.internal_routes import InternalRoutes
|
from api_server.routes.internal.internal_routes import InternalRoutes
|
||||||
from protocol import BinaryEventTypes
|
from protocol import BinaryEventTypes
|
||||||
|
|
||||||
# Time in seconds
|
# Import cache control middleware
|
||||||
ONE_HOUR: int = 3600
|
from middleware.cache_middleware import cache_control, ONE_HOUR, ONE_DAY, IMG_EXTENSIONS
|
||||||
ONE_DAY: int = 86400
|
|
||||||
IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')
|
|
||||||
|
|
||||||
async def send_socket_catch_exception(function, message):
|
async def send_socket_catch_exception(function, message):
|
||||||
try:
|
try:
|
||||||
@ -51,23 +49,6 @@ async def send_socket_catch_exception(function, message):
|
|||||||
logging.warning("send error: {}".format(err))
|
logging.warning("send error: {}".format(err))
|
||||||
|
|
||||||
@web.middleware
|
@web.middleware
|
||||||
async def cache_control(request: web.Request, handler: Callable[[web.Request], Awaitable[web.Response]]) -> web.Response:
|
|
||||||
response: web.Response = await handler(request)
|
|
||||||
|
|
||||||
if request.path.endswith('.js') or request.path.endswith('.css') or request.path.endswith('index.json'):
|
|
||||||
response.headers.setdefault('Cache-Control', 'no-cache')
|
|
||||||
elif request.path.lower().endswith(IMG_EXTENSIONS):
|
|
||||||
if response.status == 404:
|
|
||||||
response.headers.setdefault('Cache-Control', f"public, max-age={ONE_HOUR}")
|
|
||||||
elif response.status in (200, 201, 202, 203, 204, 205, 206, 301, 308):
|
|
||||||
# Success responses and permanent redirects - cache for 1 day
|
|
||||||
response.headers.setdefault('Cache-Control', f"public, max-age={ONE_DAY}")
|
|
||||||
elif response.status in (302, 303, 307):
|
|
||||||
# Temporary redirects - no cache
|
|
||||||
response.headers.setdefault('Cache-Control', 'no-cache')
|
|
||||||
# Note: 304 Not Modified falls through - no cache headers set
|
|
||||||
|
|
||||||
return response
|
|
||||||
|
|
||||||
@web.middleware
|
@web.middleware
|
||||||
async def compress_body(request: web.Request, handler):
|
async def compress_body(request: web.Request, handler):
|
||||||
|
|||||||
@ -3,20 +3,11 @@
|
|||||||
import pytest
|
import pytest
|
||||||
from aiohttp import web
|
from aiohttp import web
|
||||||
from aiohttp.test_utils import make_mocked_request
|
from aiohttp.test_utils import make_mocked_request
|
||||||
from unittest.mock import patch
|
|
||||||
|
from middleware.cache_middleware import cache_control, ONE_HOUR, ONE_DAY, IMG_EXTENSIONS
|
||||||
|
|
||||||
pytestmark = pytest.mark.asyncio # Apply asyncio mark to all tests
|
pytestmark = pytest.mark.asyncio # Apply asyncio mark to all tests
|
||||||
|
|
||||||
# Mock the problematic imports before importing server
|
|
||||||
with (
|
|
||||||
patch("app.frontend_management.FrontendManager"),
|
|
||||||
patch("utils.install_util.get_missing_requirements_message"),
|
|
||||||
patch("utils.install_util.requirements_path"),
|
|
||||||
patch("comfy.model_management.get_torch_device"),
|
|
||||||
patch("comfy.model_management.get_total_memory", return_value=1024 * 1024 * 1024),
|
|
||||||
):
|
|
||||||
from server import cache_control, ONE_HOUR, ONE_DAY, IMG_EXTENSIONS
|
|
||||||
|
|
||||||
|
|
||||||
class TestCacheControl:
|
class TestCacheControl:
|
||||||
"""Test cache control middleware functionality"""
|
"""Test cache control middleware functionality"""
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user