Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions server.py
Original file line number Diff line number Diff line change
Expand Up @@ -555,7 +555,7 @@ async def view_image(request):
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'
Expand All @@ -575,7 +575,7 @@ async def view_image(request):
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:
Expand All @@ -592,7 +592,7 @@ async def view_image(request):
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.
Expand All @@ -609,7 +609,7 @@ async def view_image(request):
return web.FileResponse(
file,
headers={
"Content-Disposition": f"filename=\"{filename}\"",
"Content-Disposition": f"attachment; filename=\"{filename}\"",
"Content-Type": content_type
}
)
Expand Down
156 changes: 156 additions & 0 deletions tests/test_view_content_disposition.py
Original file line number Diff line number Diff line change
@@ -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
Comment on lines +63 to +77
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Fail fast when server readiness never reaches 200.

If no attempt gets HTTP 200 (but no ConnectionError is thrown), the fixture still returns a session and tests fail later with less-clear errors.

Proposed fix
 `@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:
+                session.close()
                 raise
-    return session
+    else:
+        session.close()
+        pytest.fail(f"Server at {server_url} did not become ready after {n_tries} attempts")
+
+    try:
+        yield session
+    finally:
+        session.close()
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@tests/test_view_content_disposition.py` around lines 63 - 77, The client
fixture retries GET /system_stats but currently returns the session even if no
attempt returned HTTP 200 (only ConnectionError triggers raise); update the
logic in client to fail fast when readiness never reaches 200 by keeping track
of the last response (resp) or a success flag inside the retry loop and after
the loop raise an exception (or call pytest.fail) if no resp.status_code == 200
was observed; reference the client function, n_tries, resp, session and the
requests.exceptions.ConnectionError handling so the test stops early with a
clear error instead of returning a 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)