diff --git a/app/prompt_worker.py b/app/prompt_worker.py
new file mode 100644
index 000000000..51fa45ca5
--- /dev/null
+++ b/app/prompt_worker.py
@@ -0,0 +1,112 @@
+import gc
+import logging
+import time
+
+import comfy.model_management
+import execution
+import hook_breaker_ac10a0
+from comfy.cli_args import args
+
+
+def prompt_worker(q, server_instance, asset_manager):
+    current_time: float = 0.0
+    cache_ram = 0
+    cache_ram_inactive = 0
+    if not args.cache_classic and not args.cache_none and args.cache_lru <= 0:
+        cache_ram = min(10.0, max(2.0, comfy.model_management.total_ram * 0.10 / 1024.0))
+        cache_ram_inactive = min(128.0, comfy.model_management.total_ram / 1024.0)
+        if len(args.cache_ram) > 0:
+            cache_ram = args.cache_ram[0]
+        if len(args.cache_ram) > 1:
+            cache_ram_inactive = args.cache_ram[1]
+
+    cache_type = execution.CacheType.RAM_PRESSURE
+    if args.cache_classic:
+        cache_type = execution.CacheType.CLASSIC
+    elif args.cache_lru > 0:
+        cache_type = execution.CacheType.LRU
+    elif args.cache_none:
+        cache_type = execution.CacheType.NONE
+
+    e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args={ "lru" : args.cache_lru, "ram" : cache_ram, "ram_inactive" : cache_ram_inactive }, asset_manager=asset_manager )
+    last_gc_collect = 0
+    need_gc = False
+    gc_collect_interval = 10.0
+
+    while True:
+        background_scan_paused = False
+        try:
+            timeout = 1000.0
+            if need_gc:
+                timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0)
+
+            queue_item = q.get(timeout=timeout)
+            if queue_item is not None:
+                item, item_id = queue_item
+                execution_start_time = time.perf_counter()
+                prompt_id = item[1]
+                server_instance.last_prompt_id = prompt_id
+
+                sensitive = item[5]
+                extra_data = item[3].copy()
+                for k in sensitive:
+                    extra_data[k] = sensitive[k]
+
+                asset_manager.pause_background_scan()
+                background_scan_paused = True
+                e.execute(item[2], prompt_id, extra_data, item[4])
+
+                need_gc = True
+
+                remove_sensitive = lambda prompt: prompt[:5] + prompt[6:]
+                q.task_done(item_id,
+                            e.history_result,
+                            status=execution.PromptQueue.ExecutionStatus(
+                                status_str='success' if e.success else 'error',
+                                completed=e.success,
+                                messages=e.status_messages), process_item=remove_sensitive)
+                if server_instance.client_id is not None:
+                    server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server_instance.client_id)
+
+                current_time = time.perf_counter()
+                execution_time = current_time - execution_start_time
+
+                # Log Time in a more readable way after 10 minutes
+                if execution_time > 600:
+                    execution_time = time.strftime("%H:%M:%S", time.gmtime(execution_time))
+                    logging.info(f"Prompt executed in {execution_time}", extra={'color': 'green'})
+                else:
+                    logging.info("Prompt executed in {:.2f} seconds".format(execution_time), extra={'color': 'green'})
+
+            flags = q.get_flags()
+            free_memory = flags.get("free_memory", False)
+
+            if flags.get("unload_models", free_memory):
+                comfy.model_management.unload_all_models()
+                need_gc = True
+                last_gc_collect = 0
+
+            if free_memory:
+                e.reset()
+                need_gc = True
+                last_gc_collect = 0
+
+            if need_gc:
+                current_time = time.perf_counter()
+                if (current_time - last_gc_collect) > gc_collect_interval:
+                    gc.collect()
+                    comfy.model_management.soft_empty_cache()
+                    last_gc_collect = current_time
+                    need_gc = False
+                    hook_breaker_ac10a0.restore_functions()
+
+                    asset_manager.queue_output_scan()
+                    asset_manager.resume_background_scan()
+                    background_scan_paused = False
+        except BaseException:
+            if background_scan_paused:
+                try:
+                    asset_manager.resume_background_scan()
+                except Exception:
+                    logging.exception("Failed to resume background asset scanning after prompt worker failure")
+            raise
diff --git a/main.py b/main.py
index b2f96bd2e..5ec031009 100644
--- a/main.py
+++ b/main.py
@@ -240,7 +240,6 @@ execute_prestartup_script()
 # Main code
 import asyncio
 import threading
-import gc
 
 if 'torch' in sys.modules:
     logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.")
@@ -248,7 +247,7 @@ if 'torch' in sys.modules:
 
 import comfy.utils
 
-import execution
+from app.prompt_worker import prompt_worker
 import server
 from protocol import BinaryEventTypes
 import nodes
@@ -316,110 +315,6 @@ def cuda_malloc_warning():
             logging.warning("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n")
 
 
-def prompt_worker(q, server_instance, asset_manager):
-    current_time: float = 0.0
-    cache_ram = 0
-    cache_ram_inactive = 0
-    if not args.cache_classic and not args.cache_none and args.cache_lru <= 0:
-        cache_ram = min(10.0, max(2.0, comfy.model_management.total_ram * 0.10 / 1024.0))
-        cache_ram_inactive = min(128.0, comfy.model_management.total_ram / 1024.0)
-        if len(args.cache_ram) > 0:
-            cache_ram = args.cache_ram[0]
-        if len(args.cache_ram) > 1:
-            cache_ram_inactive = args.cache_ram[1]
-
-    cache_type = execution.CacheType.RAM_PRESSURE
-    if args.cache_classic:
-        cache_type = execution.CacheType.CLASSIC
-    elif args.cache_lru > 0:
-        cache_type = execution.CacheType.LRU
-    elif args.cache_none:
-        cache_type = execution.CacheType.NONE
-
-    e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args={ "lru" : args.cache_lru, "ram" : cache_ram, "ram_inactive" : cache_ram_inactive }, asset_manager=asset_manager )
-    last_gc_collect = 0
-    need_gc = False
-    gc_collect_interval = 10.0
-
-    while True:
-        background_scan_paused = False
-        try:
-            timeout = 1000.0
-            if need_gc:
-                timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0)
-
-            queue_item = q.get(timeout=timeout)
-            if queue_item is not None:
-                item, item_id = queue_item
-                execution_start_time = time.perf_counter()
-                prompt_id = item[1]
-                server_instance.last_prompt_id = prompt_id
-
-                sensitive = item[5]
-                extra_data = item[3].copy()
-                for k in sensitive:
-                    extra_data[k] = sensitive[k]
-
-                asset_manager.pause_background_scan()
-                background_scan_paused = True
-                e.execute(item[2], prompt_id, extra_data, item[4])
-
-                need_gc = True
-
-                remove_sensitive = lambda prompt: prompt[:5] + prompt[6:]
-                q.task_done(item_id,
-                            e.history_result,
-                            status=execution.PromptQueue.ExecutionStatus(
-                                status_str='success' if e.success else 'error',
-                                completed=e.success,
-                                messages=e.status_messages), process_item=remove_sensitive)
-                if server_instance.client_id is not None:
-                    server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server_instance.client_id)
-
-                current_time = time.perf_counter()
-                execution_time = current_time - execution_start_time
-
-                # Log Time in a more readable way after 10 minutes
-                if execution_time > 600:
-                    execution_time = time.strftime("%H:%M:%S", time.gmtime(execution_time))
-                    logging.info(f"Prompt executed in {execution_time}", extra={'color': 'green'})
-                else:
-                    logging.info("Prompt executed in {:.2f} seconds".format(execution_time), extra={'color': 'green'})
-
-            flags = q.get_flags()
-            free_memory = flags.get("free_memory", False)
-
-            if flags.get("unload_models", free_memory):
-                comfy.model_management.unload_all_models()
-                need_gc = True
-                last_gc_collect = 0
-
-            if free_memory:
-                e.reset()
-                need_gc = True
-                last_gc_collect = 0
-
-            if need_gc:
-                current_time = time.perf_counter()
-                if (current_time - last_gc_collect) > gc_collect_interval:
-                    gc.collect()
-                    comfy.model_management.soft_empty_cache()
-                    last_gc_collect = current_time
-                    need_gc = False
-                    hook_breaker_ac10a0.restore_functions()
-
-                    asset_manager.queue_output_scan()
-                    asset_manager.resume_background_scan()
-                    background_scan_paused = False
-        except BaseException:
-            if background_scan_paused:
-                try:
-                    asset_manager.resume_background_scan()
-                except Exception:
-                    logging.exception("Failed to resume background asset scanning after prompt worker failure")
-            raise
-
-
 async def run(server_instance, address='', port=8188, verbose=True, call_on_start=None):
     addresses = []
     for addr in address.split(","):
diff --git a/tests-unit/app_test/test_prompt_worker.py b/tests-unit/app_test/test_prompt_worker.py
index 3f96b50bf..0634beb06 100644
--- a/tests-unit/app_test/test_prompt_worker.py
+++ b/tests-unit/app_test/test_prompt_worker.py
@@ -1,170 +1,102 @@
-import subprocess
-import sys
-from pathlib import Path
+import importlib
+
+import pytest
 
 
-def _run_prompt_worker(script: str) -> subprocess.CompletedProcess[str]:
-    return subprocess.run(
-        [sys.executable, "-c", script],
-        cwd=Path(__file__).parents[2],
-        capture_output=True,
-        text=True,
-        check=False,
-    )
+class LoopEscape(Exception):
+    pass
 
 
-def test_prompt_worker_resumes_background_scan_when_execute_raises() -> None:
-    script = """
-import sys
-
-sys.argv = ["main.py", "--cpu"]
-
-import main
-
 class Queue:
+    def __init__(self, completion_error: RuntimeError | None = None) -> None:
+        self.completion_error = completion_error
+        self.get_calls = 0
+
     def get(self, timeout=None):
+        self.get_calls += 1
+        if self.get_calls > 1:
+            raise LoopEscape("prompt worker requested a second item")
         return (0, "prompt-id", {}, {}, [], {}), 1
 
+    def task_done(self, *args, **kwargs) -> None:
+        if self.completion_error is not None:
+            raise self.completion_error
+
+
 class Server:
     last_prompt_id = None
     client_id = None
 
-class BackgroundScan:
-    paused = False
 
-    def pause_background_scan(self):
+class AssetManager:
+    def __init__(self, resume_error: RuntimeError | None = None) -> None:
+        self.paused = False
+        self.resume_error = resume_error
+
+    def pause_background_scan(self) -> None:
         self.paused = True
 
-    def resume_background_scan(self):
+    def resume_background_scan(self) -> None:
         self.paused = False
+        if self.resume_error is not None:
+            raise self.resume_error
+
 
 class Executor:
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         self.history_result = {}
         self.success = True
         self.status_messages = []
 
-    def execute(self, *args, **kwargs):
-        raise RuntimeError("forced execute failure")
-
-main.args.cache_classic = True
-main.execution.PromptExecutor = Executor
-asset_manager = BackgroundScan()
-
-try:
-    main.prompt_worker(Queue(), Server(), asset_manager)
-except RuntimeError as error:
-    assert str(error) == "forced execute failure"
-
-assert asset_manager.paused is False
-"""
-
-    result = _run_prompt_worker(script)
-
-    assert result.returncode == 0, result.stderr
-
-
-def test_prompt_worker_resumes_background_scan_when_completion_raises() -> None:
-    script = """
-import sys
-
-sys.argv = ["main.py", "--cpu"]
-
-import main
-
-class Queue:
-    def get(self, timeout=None):
-        return (0, "prompt-id", {}, {}, [], {}), 1
-
-    def task_done(self, *args, **kwargs):
-        raise RuntimeError("forced completion failure")
-
-class Server:
-    last_prompt_id = None
-    client_id = None
-
-class BackgroundScan:
-    paused = False
-
-    def pause_background_scan(self):
-        self.paused = True
-
-    def resume_background_scan(self):
-        self.paused = False
-
-class Executor:
-    def __init__(self, *args, **kwargs):
-        self.history_result = {}
-        self.success = True
-        self.status_messages = []
-
-    def execute(self, *args, **kwargs):
+    def execute(self, *args, **kwargs) -> None:
         return None
 
-main.args.cache_classic = True
-main.execution.PromptExecutor = Executor
-asset_manager = BackgroundScan()
 
-try:
-    main.prompt_worker(Queue(), Server(), asset_manager)
-except RuntimeError as error:
-    assert str(error) == "forced completion failure"
-
-assert asset_manager.paused is False
-"""
-
-    result = _run_prompt_worker(script)
-
-    assert result.returncode == 0, result.stderr
-
-
-def test_prompt_worker_preserves_execute_error_when_resume_raises() -> None:
-    script = """
-import sys
-
-sys.argv = ["main.py", "--cpu"]
-
-import main
-
-class Queue:
-    def get(self, timeout=None):
-        return (0, "prompt-id", {}, {}, [], {}), 1
-
-class Server:
-    last_prompt_id = None
-    client_id = None
-
-class BackgroundScan:
-    paused = False
-
-    def pause_background_scan(self):
-        self.paused = True
-
-    def resume_background_scan(self):
-        self.paused = False
-        raise RuntimeError("forced resume failure")
-
-class Executor:
-    def __init__(self, *args, **kwargs):
-        pass
-
-    def execute(self, *args, **kwargs):
+class ExecuteFailureExecutor(Executor):
+    def execute(self, *args, **kwargs) -> None:
         raise RuntimeError("forced execute failure")
 
-main.args.cache_classic = True
-main.execution.PromptExecutor = Executor
-asset_manager = BackgroundScan()
 
-try:
-    main.prompt_worker(Queue(), Server(), asset_manager)
-except RuntimeError as error:
-    assert str(error) == "forced execute failure"
-else:
-    raise AssertionError("prompt_worker should propagate the execute failure")
+@pytest.fixture
+def prompt_worker_module(monkeypatch):
+    from comfy.cli_args import args
 
-assert asset_manager.paused is False
-"""
+    monkeypatch.setattr(args, "cpu", True, raising=False)
+    try:
+        return importlib.import_module("app.prompt_worker")
+    except Exception as exc:
+        pytest.skip(f"prompt worker module could not be imported in CPU mode: {exc!r}")
 
-    result = _run_prompt_worker(script)
 
-    assert result.returncode == 0, result.stderr
+def test_prompt_worker_resumes_background_scan_when_execute_raises(prompt_worker_module, monkeypatch) -> None:
+    monkeypatch.setattr(prompt_worker_module.execution, "PromptExecutor", ExecuteFailureExecutor)
+    asset_manager = AssetManager()
+
+    with pytest.raises(RuntimeError, match="^forced execute failure$"):
+        prompt_worker_module.prompt_worker(Queue(), Server(), asset_manager)
+
+    assert asset_manager.paused is False
+
+
+def test_prompt_worker_resumes_background_scan_when_completion_raises(prompt_worker_module, monkeypatch) -> None:
+    monkeypatch.setattr(prompt_worker_module.execution, "PromptExecutor", Executor)
+    asset_manager = AssetManager()
+
+    with pytest.raises(RuntimeError, match="^forced completion failure$"):
+        prompt_worker_module.prompt_worker(
+            Queue(completion_error=RuntimeError("forced completion failure")),
+            Server(),
+            asset_manager,
+        )
+
+    assert asset_manager.paused is False
+
+
+def test_prompt_worker_preserves_execute_error_when_resume_raises(prompt_worker_module, monkeypatch) -> None:
+    monkeypatch.setattr(prompt_worker_module.execution, "PromptExecutor", ExecuteFailureExecutor)
+    asset_manager = AssetManager(resume_error=RuntimeError("forced resume failure"))
+
+    with pytest.raises(RuntimeError, match="^forced execute failure$"):
+        prompt_worker_module.prompt_worker(Queue(), Server(), asset_manager)
+
+    assert asset_manager.paused is False