feat(middleware): add cache control middleware to manage response caching for image files

middleware/__init__.py

@@ -0,0 +1 @@
+"""Server middleware modules"""

middleware/cache_middleware.py

@@ -0,0 +1,52 @@
+"""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")
+        return response
+
+    # Early return for non-image files - no cache headers needed
+    if not request.path.lower().endswith(IMG_EXTENSIONS):
+        return response
+
+    # Handle image files
+    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

standalone.py

@@ -2,6 +2,7 @@ from pathlib import Path
 import os
 import sys
 import json
+from middleware.cache_middleware import cache_control
 
 # Create mock modules for py/nodes directory - add this before any other imports
 def mock_nodes_directory():
@@ -129,7 +130,7 @@ class StandaloneServer:
     """Server implementation for standalone mode"""
     
     def __init__(self):
-        self.app = web.Application(logger=logger)
+        self.app = web.Application(logger=logger, middlewares=[cache_control])
         self.instance = self  # Make it compatible with PromptServer.instance pattern
         
         # Ensure the app's access logger is configured to reduce verbosity