From 30956aeefc4d394a70de7f0385e5a37ac4f9c230 Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Tue, 16 Sep 2025 15:05:31 +0800
Subject: [PATCH] feat(middleware): add cache control middleware to manage
 response caching for image files

---
 middleware/__init__.py         |  1 +
 middleware/cache_middleware.py | 52 ++++++++++++++++++++++++++++++++++
 standalone.py                  |  3 +-
 3 files changed, 55 insertions(+), 1 deletion(-)
 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 00000000..2d7c7c3a
--- /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 00000000..374ef793
--- /dev/null
+++ b/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
diff --git a/standalone.py b/standalone.py
index b8532abf..56d4545d 100644
--- a/standalone.py
+++ b/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