diff --git a/glob/manager_server.py b/glob/manager_server.py index 730cc18d..591509b9 100644 --- a/glob/manager_server.py +++ b/glob/manager_server.py @@ -1967,17 +1967,7 @@ def restart(self): print("\nRestarting... [Legacy Mode]\n\n") # This printing should not be logging - that will be ugly - sys_argv = sys.argv.copy() - if '--windows-standalone-build' in sys_argv: - sys_argv.remove('--windows-standalone-build') - - if sys_argv[0].endswith("__main__.py"): # this is a python module - module_name = os.path.basename(os.path.dirname(sys_argv[0])) - cmds = [sys.executable, '-m', module_name] + sys_argv[1:] - elif sys.platform.startswith('win32'): - cmds = ['"' + sys.executable + '"', '"' + sys_argv[0] + '"'] + sys_argv[1:] - else: - cmds = [sys.executable] + sys_argv + cmds = manager_util.build_restart_command(sys.argv, drop_args=('--windows-standalone-build',)) print(f"Command: {cmds}", flush=True) diff --git a/glob/manager_util.py b/glob/manager_util.py index 6bf1a055..0bbb5c32 100644 --- a/glob/manager_util.py +++ b/glob/manager_util.py @@ -693,6 +693,32 @@ def robust_readlines(fullpath): return [] +def build_restart_command(argv, drop_args=()): + """ + Rebuild the command line used to re-exec the current ComfyUI process. + + `argv` is `sys.argv` of the running process; entries listed in `drop_args` are left out. + The result is meant to be passed to `os.execv(sys.executable, ...)`. + + On Windows, `os.execv` does not quote anything: it joins the list with spaces into one + command line and the new process splits it again. Every entry is therefore quoted with + the MSVC rules (`subprocess.list2cmdline`), otherwise an argument such as `Profile 1` + comes back as two arguments. Other platforms pass the list through untouched. + """ + argv = [x for x in argv if x not in drop_args] + + if argv[0].endswith("__main__.py"): # this is a python module + module_name = os.path.basename(os.path.dirname(argv[0])) + cmds = [sys.executable, '-m', module_name] + argv[1:] + else: + cmds = [sys.executable] + argv + + if sys.platform.startswith('win32'): + cmds = [subprocess.list2cmdline([x]) for x in cmds] + + return cmds + + def restore_pip_snapshot(pips, options): non_url = [] local_url = [] diff --git a/prestartup_script.py b/prestartup_script.py index 061e448b..da707f88 100644 --- a/prestartup_script.py +++ b/prestartup_script.py @@ -864,15 +864,7 @@ if script_executed: print("--------------------------------------------------------------------------\n") exit(0) else: - sys_argv = sys.argv.copy() - - if sys_argv[0].endswith("__main__.py"): # this is a python module - module_name = os.path.basename(os.path.dirname(sys_argv[0])) - cmds = [sys.executable, '-m', module_name] + sys_argv[1:] - elif sys.platform.startswith('win32'): - cmds = ['"' + sys.executable + '"', '"' + sys_argv[0] + '"'] + sys_argv[1:] - else: - cmds = [sys.executable] + sys_argv + cmds = manager_util.build_restart_command(sys.argv) print(f"Command: {cmds}", flush=True) print("--------------------------------------------------------------------------\n") diff --git a/tests/test_restart_command.py b/tests/test_restart_command.py new file mode 100644 index 00000000..a71a254c --- /dev/null +++ b/tests/test_restart_command.py @@ -0,0 +1,154 @@ +"""Regression tests for the restart command line (#3339). + +Both restart paths (the reboot endpoint and the prestartup restart after a +dependency install) re-exec ComfyUI with os.execv. On Windows os.execv joins +its arguments with spaces and quotes nothing, so `--browser-profile "Profile 1"` +used to reach the new process as `--browser-profile Profile 1`. +""" +import ast +import json +import os +import subprocess +import sys +import time + +import pytest + +from manager_test_utils import GLOB_DIR, load_manager_util + +MANAGER_UTIL = load_manager_util() + +BACKSLASH = chr(92) +TRICKY_ARGS = [ + "--browser-profile", "Profile 1", + "--output-directory", "C:" + BACKSLASH + "my dir" + BACKSLASH, + "--note", 'say "hi"', + "--empty", "", + "--port", "8188", +] + + +def test_windows_quotes_only_what_needs_quoting(monkeypatch): + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(sys, "executable", "C:" + BACKSLASH + "Program Files" + BACKSLASH + "python.exe") + + cmds = MANAGER_UTIL.build_restart_command(["C:" + BACKSLASH + "Comfy UI" + BACKSLASH + "main.py"] + TRICKY_ARGS) + + assert cmds == [ + '"C:' + BACKSLASH + 'Program Files' + BACKSLASH + 'python.exe"', + '"C:' + BACKSLASH + 'Comfy UI' + BACKSLASH + 'main.py"', + "--browser-profile", '"Profile 1"', + "--output-directory", '"C:' + BACKSLASH + 'my dir' + BACKSLASH * 2 + '"', + "--note", '"say ' + BACKSLASH + '"hi' + BACKSLASH + '""', + "--empty", '""', + "--port", "8188", + ] + + +def test_posix_passes_arguments_through(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + monkeypatch.setattr(sys, "executable", "/usr/bin/python3") + + argv = ["/opt/Comfy UI/main.py"] + TRICKY_ARGS + + assert MANAGER_UTIL.build_restart_command(argv) == ["/usr/bin/python3"] + argv + + +def test_module_entry_point_uses_dash_m(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + monkeypatch.setattr(sys, "executable", "/usr/bin/python3") + + cmds = MANAGER_UTIL.build_restart_command(["/opt/site-packages/comfy/__main__.py", "--port", "8188"]) + + assert cmds == ["/usr/bin/python3", "-m", "comfy", "--port", "8188"] + + +def test_module_entry_point_is_quoted_on_windows(monkeypatch): + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(sys, "executable", "C:" + BACKSLASH + "Program Files" + BACKSLASH + "python.exe") + + # Forward slashes: the module name is then found by both the POSIX and the Windows os.path + cmds = MANAGER_UTIL.build_restart_command( + ["C:/site packages/comfy/__main__.py", "--browser-profile", "Profile 1"]) + + assert cmds == [ + '"C:' + BACKSLASH + 'Program Files' + BACKSLASH + 'python.exe"', + "-m", "comfy", + "--browser-profile", '"Profile 1"', + ] + + +def test_drop_args_removes_only_listed_arguments(monkeypatch): + monkeypatch.setattr(sys, "platform", "linux") + monkeypatch.setattr(sys, "executable", "/usr/bin/python3") + argv = ["main.py", "--windows-standalone-build", "--cpu"] + + assert MANAGER_UTIL.build_restart_command(argv) == ["/usr/bin/python3"] + argv + assert MANAGER_UTIL.build_restart_command(argv, drop_args=("--windows-standalone-build",)) == \ + ["/usr/bin/python3", "main.py", "--cpu"] + assert argv == ["main.py", "--windows-standalone-build", "--cpu"] # caller's list is untouched + + +def restart_call_drop_args(path): + """The `drop_args` that the one restart site in `path` passes for `sys.argv`.""" + calls = [node for node in ast.walk(ast.parse(path.read_text(encoding="utf-8"))) + if isinstance(node, ast.Call) and ast.unparse(node.func) == "manager_util.build_restart_command"] + assert len(calls) == 1, f"expected one restart site in {path.name}" + call = calls[0] + assert ast.unparse(call.args[0]) == "sys.argv" + + drop_args = call.args[1:] + [keyword.value for keyword in call.keywords if keyword.arg == "drop_args"] + assert len(call.args) + len(call.keywords) == 1 + len(drop_args), "unexpected argument" + return tuple(ast.literal_eval(drop_args[0])) if drop_args else () + + +def test_only_the_reboot_endpoint_drops_the_standalone_build_flag(): + # ComfyUI opens a browser when it starts with this flag. The reboot endpoint restarts a + # ComfyUI whose page is already open, so it drops the flag (#1160). + assert restart_call_drop_args(GLOB_DIR / "manager_server.py") == ("--windows-standalone-build",) + + prestartup_script = GLOB_DIR.resolve().parent / "prestartup_script.py" + if os.environ.get("MANAGER_GLOB_DIR") and not prestartup_script.exists(): + pytest.skip("MANAGER_GLOB_DIR points at a glob directory without prestartup_script.py beside it") + # The prestartup restart happens before ComfyUI has opened anything, so the flag must stay. + assert restart_call_drop_args(prestartup_script) == () + + +LAUNCHER = """ +import importlib.util, os, sys +spec = importlib.util.spec_from_file_location("_manager_util_under_test", sys.argv[1]) +manager_util = importlib.util.module_from_spec(spec) +spec.loader.exec_module(manager_util) +os.execv(sys.executable, manager_util.build_restart_command(sys.argv[2:])) +""" + +CHILD = """ +import json, sys +with open(sys.argv[1] + ".tmp", "w", encoding="utf-8") as f: + json.dump(sys.argv[2:], f) +import os +os.replace(sys.argv[1] + ".tmp", sys.argv[1]) +""" + + +def test_arguments_survive_a_real_execv(tmp_path): + """End to end on the running platform: what the re-exec'ed process sees must equal what was sent.""" + work_dir = tmp_path / "dir with spaces" + work_dir.mkdir() + child = work_dir / "main entry.py" + child.write_text(CHILD, encoding="utf-8") + result_file = work_dir / "argv seen.json" + + subprocess.run( + [sys.executable, "-c", LAUNCHER, str(GLOB_DIR / "manager_util.py"), str(child), str(result_file)] + TRICKY_ARGS, + check=True, timeout=60, + ) + + # On Windows os.execv returns control to the caller before the new process is done. + deadline = time.time() + 30 + while not result_file.exists(): + if time.time() > deadline: + pytest.fail("re-exec'ed process never reported its arguments") + time.sleep(0.1) + + assert json.loads(result_file.read_text(encoding="utf-8")) == TRICKY_ARGS