Merge 1f8a74b76f into b353a7c863

server.py

@@ -555,7 +555,7 @@ class PromptServer():
                             buffer.seek(0)
 
                             return web.Response(body=buffer.read(), content_type=f'image/{image_format}',
-                                                headers={"Content-Disposition": f"filename=\"{filename}\""})
+                                                headers={"Content-Disposition": f"attachment; filename=\"{filename}\""})
 
                     if 'channel' not in request.rel_url.query:
                         channel = 'rgba'
@@ -575,7 +575,7 @@ class PromptServer():
                             buffer.seek(0)
 
                             return web.Response(body=buffer.read(), content_type='image/png',
-                                                headers={"Content-Disposition": f"filename=\"{filename}\""})
+                                                headers={"Content-Disposition": f"attachment; filename=\"{filename}\""})
 
                     elif channel == 'a':
                         with Image.open(file) as img:
@@ -592,7 +592,7 @@ class PromptServer():
                             alpha_buffer.seek(0)
 
                             return web.Response(body=alpha_buffer.read(), content_type='image/png',
-                                                headers={"Content-Disposition": f"filename=\"{filename}\""})
+                                                headers={"Content-Disposition": f"attachment; filename=\"{filename}\""})
                     else:
                         # Use the content type from asset resolution if available,
                         # otherwise guess from the filename.
@@ -609,7 +609,7 @@ class PromptServer():
                         return web.FileResponse(
                             file,
                             headers={
-                                "Content-Disposition": f"filename=\"{filename}\"",
+                                "Content-Disposition": f"attachment; filename=\"{filename}\"",
                                 "Content-Type": content_type
                             }
                         )

tests/test_view_content_disposition.py

@@ -0,0 +1,156 @@
+"""
+Tests for the /view endpoint Content-Disposition header fix.
+
+Verifies that the Content-Disposition header complies with RFC 2183
+by using `attachment; filename="..."` format instead of just `filename="..."`.
+
+See: https://github.com/Comfy-Org/ComfyUI/issues/8914
+"""
+
+import pytest
+import subprocess
+import time
+import os
+import tempfile
+import struct
+import pytest
+
+# Simple 1x1 PNG image (minimal valid PNG)
+MINIMAL_PNG = bytes([
+    0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A,  # PNG signature
+    0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52,  # IHDR chunk length + type
+    0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,  # 1x1 dimensions
+    0x08, 0x02, 0x00, 0x00, 0x00,                     # bit depth, color type, etc
+    0x90, 0x77, 0x53, 0xDE,                           # IHDR CRC
+    0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, 0x54,  # IDAT chunk
+    0x08, 0xD7, 0x63, 0xF8, 0xFF, 0xFF, 0xFF, 0x00,  # image data
+    0x05, 0xFE, 0x02, 0xFE,                           # IDAT CRC
+    0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44,  # IEND chunk
+    0xAE, 0x42, 0x60, 0x82,                           # IEND CRC
+])
+
+
+@pytest.mark.execution
+class TestViewContentDisposition:
+    """Test suite for Content-Disposition header in /view endpoint."""
+
+    @pytest.fixture(scope="class", autouse=True)
+    def output_dir(self, args_pytest):
+        """Use the test output directory."""
+        return args_pytest["output_dir"]
+
+    @pytest.fixture(scope="class", autouse=True)
+    def _server(self, args_pytest):
+        """Start ComfyUI server for testing."""
+        pargs = [
+            'python', 'main.py',
+            '--output-directory', args_pytest["output_dir"],
+            '--listen', args_pytest["listen"],
+            '--port', str(args_pytest["port"]),
+            '--cpu',
+        ]
+        p = subprocess.Popen(pargs, cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+        yield
+        p.kill()
+        p.wait()
+
+    @pytest.fixture(scope="class")
+    def server_url(self, args_pytest):
+        """Get the server base URL."""
+        return f"http://{args_pytest['listen']}:{args_pytest['port']}"
+
+    @pytest.fixture(scope="class")
+    def client(self, server_url, _server):
+        """Create HTTP client with retry."""
+        import requests
+        session = requests.Session()
+        n_tries = 10
+        for i in range(n_tries):
+            time.sleep(2)
+            try:
+                resp = session.get(f"{server_url}/system_stats", timeout=5)
+                if resp.status_code == 200:
+                    break
+            except requests.exceptions.ConnectionError:
+                if i == n_tries - 1:
+                    raise
+        return session
+
+    @pytest.fixture
+    def test_image_file(self, output_dir):
+        """Create a temporary test image file and clean it up after the test."""
+        filename = "test_content_disposition.png"
+        filepath = os.path.join(output_dir, filename)
+        with open(filepath, 'wb') as f:
+            f.write(MINIMAL_PNG)
+        yield filename, filepath
+        # Cleanup
+        if os.path.exists(filepath):
+            os.remove(filepath)
+
+    def test_view_content_disposition_is_attachment_format(self, client, server_url, test_image_file):
+        """Test that /view endpoint returns Content-Disposition with RFC 2183 compliant format.
+
+        The header should be: Content-Disposition: attachment; filename="test_content_disposition.png"
+        NOT: Content-Disposition: filename="test_content_disposition.png"
+
+        This test validates the fix for: https://github.com/Comfy-Org/ComfyUI/issues/8914
+        """
+        filename, _ = test_image_file
+        response = client.get(
+            f"{server_url}/view",
+            params={"filename": filename},
+            timeout=10
+        )
+
+        assert response.status_code == 200, f"Expected 200, got {response.status_code}"
+
+        content_disposition = response.headers.get("Content-Disposition", "")
+        assert content_disposition, "Content-Disposition header should be present"
+
+        # RFC 2183 compliant format: "attachment; filename=\"name.ext\""
+        expected_prefix = f"attachment; filename=\"{filename}\""
+        assert content_disposition == expected_prefix, (
+            f"Content-Disposition header should be RFC 2183 compliant.\n"
+            f"Expected: {expected_prefix}\n"
+            f"Got: {content_disposition}"
+        )
+
+    def test_view_with_subfolder_content_disposition(self, client, server_url, output_dir):
+        """Test Content-Disposition with subfolder parameter."""
+        import requests
+
+        # Create a subfolder and test image
+        subfolder = "test_subfolder"
+        subfolder_path = os.path.join(output_dir, subfolder)
+        os.makedirs(subfolder_path, exist_ok=True)
+
+        filename = "test_subfolder.png"
+        filepath = os.path.join(subfolder_path, filename)
+        with open(filepath, 'wb') as f:
+            f.write(MINIMAL_PNG)
+
+        try:
+            response = client.get(
+                f"{server_url}/view",
+                params={"filename": filename, "subfolder": subfolder},
+                timeout=10
+            )
+
+            assert response.status_code == 200, f"Expected 200, got {response.status_code}"
+
+            content_disposition = response.headers.get("Content-Disposition", "")
+            assert content_disposition, "Content-Disposition header should be present"
+
+            expected_prefix = f"attachment; filename=\"{filename}\""
+            assert content_disposition == expected_prefix, (
+                f"Content-Disposition header should be RFC 2183 compliant.\n"
+                f"Expected: {expected_prefix}\n"
+                f"Got: {content_disposition}"
+            )
+        finally:
+            # Cleanup
+            if os.path.exists(filepath):
+                os.remove(filepath)
+            if os.path.exists(subfolder_path):
+                os.rmdir(subfolder_path)