From a0e084042e63c8c560760128201b9a8a824f9402 Mon Sep 17 00:00:00 2001
From: Arjan Singh <arjan@comfy.org>
Date: Fri, 29 Aug 2025 11:27:07 -0700
Subject: [PATCH] refactor: move cache-controls to /middleware

We can test the code without excessively mocking server dependencies
---
 middleware/__init__.py                       |  1 +
 middleware/cache_middleware.py               | 46 ++++++++++++++++++++
 server.py                                    | 23 +---------
 tests-unit/server_test/test_cache_control.py | 13 +-----
 4 files changed, 51 insertions(+), 32 deletions(-)
 create mode 100644 middleware/__init__.py
 create mode 100644 middleware/cache_middleware.py

diff --git a/middleware/__init__.py b/middleware/__init__.py
new file mode 100644
index 000000000..2d7c7c3a9
--- /dev/null
+++ b/middleware/__init__.py
@@ -0,0 +1 @@
+"""Server middleware modules"""
diff --git a/middleware/cache_middleware.py b/middleware/cache_middleware.py
new file mode 100644
index 000000000..c31fdc851
--- /dev/null
+++ b/middleware/cache_middleware.py
@@ -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
diff --git a/server.py b/server.py
index 44ae24a85..df89756ff 100644
--- a/server.py
+++ b/server.py
@@ -39,10 +39,8 @@ from typing import Optional, Union, Callable, Awaitable
 from api_server.routes.internal.internal_routes import InternalRoutes
 from protocol import BinaryEventTypes
 
-# Time in seconds
-ONE_HOUR: int = 3600
-ONE_DAY: int = 86400
-IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')
+# Import cache control middleware
+from middleware.cache_middleware import cache_control, ONE_HOUR, ONE_DAY, IMG_EXTENSIONS
 
 async def send_socket_catch_exception(function, message):
     try:
@@ -51,23 +49,6 @@ async def send_socket_catch_exception(function, message):
         logging.warning("send error: {}".format(err))
 
 @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
 async def compress_body(request: web.Request, handler):
diff --git a/tests-unit/server_test/test_cache_control.py b/tests-unit/server_test/test_cache_control.py
index 183490b4d..da3968ae3 100644
--- a/tests-unit/server_test/test_cache_control.py
+++ b/tests-unit/server_test/test_cache_control.py
@@ -3,20 +3,11 @@
 import pytest
 from aiohttp import web
 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
 
-# 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:
     """Test cache control middleware functionality"""