fix(product-suite): finalize and cancel generation

tests/test_ai.py

@@ -6,6 +6,7 @@ import socket
 import sys
 import threading
 import unittest
+from concurrent.futures import CancelledError
 from types import SimpleNamespace
 from unittest import mock
 
@@ -38,6 +39,7 @@ class _RequestsResponse:
         self.content = content
         self.headers = headers or {}
         self.text = json.dumps(self.payload, ensure_ascii=False)
+        self.closed = False
 
     def json(self):
         return self.payload
@@ -46,6 +48,9 @@ class _RequestsResponse:
         if self.content:
             yield self.content
 
+    def close(self):
+        self.closed = True
+
 class AITests(TempDirMixin, unittest.TestCase):
     def _write_models(self, path, text=None, image=None):
         text = text or {
@@ -732,7 +737,16 @@ class AITests(TempDirMixin, unittest.TestCase):
         ]
         calls = []
 
-        def fake_run(args, **kwargs):
+        class FakeProcess:
+            returncode = 0
+
+            def poll(self):
+                return self.returncode
+
+            def communicate(self):
+                return b"", b""
+
+        def fake_popen(args, **kwargs):
             calls.append((args, kwargs))
             self.assertIn("-K", args)
             config_path = args[args.index("-K") + 1]
@@ -755,12 +769,12 @@ class AITests(TempDirMixin, unittest.TestCase):
             output_path = args[args.index("--output") + 1]
             with open(output_path, "wb") as fh:
                 fh.write(generated_png)
-            return SimpleNamespace(returncode=0, stdout=b"", stderr=b"")
+            return FakeProcess()
 
         with mock.patch("app.ai.os.name", "nt"), \
                 mock.patch("app.ai.subprocess.CREATE_NO_WINDOW", 0x08000000, create=True), \
                 mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
-                mock.patch("app.ai.subprocess.run", side_effect=fake_run), \
+                mock.patch("app.ai.subprocess.Popen", side_effect=fake_popen), \
                 mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
             image_bytes = ai._download_cmhub_image(
                 url,
@@ -786,7 +800,7 @@ class AITests(TempDirMixin, unittest.TestCase):
 
     def test_cmhub_image_download_skips_curl_for_private_url(self):
         with mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
-                mock.patch("app.ai.subprocess.run") as run:
+                mock.patch("app.ai.subprocess.Popen") as popen:
             with self.assertRaises(ai.AIError):
                 ai._download_cmhub_image(
                     "http://127.0.0.1/a.png",
@@ -794,7 +808,7 @@ class AITests(TempDirMixin, unittest.TestCase):
                     read_timeout=ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS,
                     download_with_curl="true",
                 )
-        run.assert_not_called()
+        popen.assert_not_called()
 
     def test_cmhub_image_download_falls_back_to_requests_when_curl_fails(self):
         generated_png = self._png_bytes()
@@ -805,11 +819,20 @@ class AITests(TempDirMixin, unittest.TestCase):
         def fake_get(url, **kwargs):
             return _RequestsResponse(content=generated_png)
 
+        class FailedProcess:
+            returncode = 28
+
+            def poll(self):
+                return self.returncode
+
+            def communicate(self):
+                return b"", b"timeout"
+
         with mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
                 mock.patch(
-                    "app.ai.subprocess.run",
-                    return_value=SimpleNamespace(returncode=28, stdout=b"", stderr=b"timeout"),
-                ) as run, \
+                    "app.ai.subprocess.Popen",
+                    return_value=FailedProcess(),
+                ) as popen, \
                 mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get) as get, \
                 mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
             image_bytes = ai._download_cmhub_image(
@@ -820,7 +843,7 @@ class AITests(TempDirMixin, unittest.TestCase):
             )
 
         self.assertEqual(generated_png, image_bytes)
-        self.assertEqual(1, run.call_count)
+        self.assertEqual(1, popen.call_count)
         self.assertEqual(1, get.call_count)
 
     def test_cmhub_image_download_auto_without_curl_uses_requests(self):
@@ -834,7 +857,7 @@ class AITests(TempDirMixin, unittest.TestCase):
 
         with mock.patch("app.ai.os.name", "nt"), \
                 mock.patch("app.ai._find_system_curl", return_value=""), \
-                mock.patch("app.ai.subprocess.run") as run, \
+                mock.patch("app.ai.subprocess.Popen") as popen, \
                 mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
                 mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
             image_bytes = ai._download_cmhub_image(
@@ -845,7 +868,7 @@ class AITests(TempDirMixin, unittest.TestCase):
             )
 
         self.assertEqual(generated_png, image_bytes)
-        run.assert_not_called()
+        popen.assert_not_called()
 
     def test_cmhub_image_download_auto_on_non_windows_uses_requests(self):
         generated_png = self._png_bytes()
@@ -858,7 +881,7 @@ class AITests(TempDirMixin, unittest.TestCase):
 
         with mock.patch("app.ai.os.name", "posix"), \
                 mock.patch("app.ai._find_system_curl", return_value="/usr/bin/curl"), \
-                mock.patch("app.ai.subprocess.run") as run, \
+                mock.patch("app.ai.subprocess.Popen") as popen, \
                 mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
                 mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
             image_bytes = ai._download_cmhub_image(
@@ -869,7 +892,93 @@ class AITests(TempDirMixin, unittest.TestCase):
             )
 
         self.assertEqual(generated_png, image_bytes)
-        run.assert_not_called()
+        popen.assert_not_called()
+
+    def test_cmhub_requests_download_cancel_closes_response(self):
+        stopped = {"value": False}
+
+        class StreamingResponse(_RequestsResponse):
+            def iter_content(self, chunk_size=65536):
+                yield b"first"
+                stopped["value"] = True
+                yield b"second"
+
+        response = StreamingResponse()
+        with mock.patch.object(
+            ai._cmhub_session(),
+            "get",
+            return_value=response,
+        ):
+            with self.assertRaises(CancelledError):
+                ai._download_cmhub_image_with_requests(
+                    "https://cdn.example.com/generated.png",
+                    connect_timeout=3,
+                    read_timeout=30,
+                    max_bytes=ai.CMHUB_IMAGE_MAX_BYTES,
+                    should_stop=lambda: stopped["value"],
+                )
+
+        self.assertTrue(response.closed)
+
+    def test_cmhub_curl_download_cancel_terminates_process_and_cleans_temp_files(self):
+        stopped = {"value": False}
+        captured_paths = []
+
+        class RunningProcess:
+            def __init__(self):
+                self.returncode = None
+                self.terminated = False
+                self.killed = False
+
+            def poll(self):
+                return self.returncode
+
+            def terminate(self):
+                self.terminated = True
+                self.returncode = -15
+
+            def kill(self):
+                self.killed = True
+                self.returncode = -9
+
+            def wait(self, timeout=None):
+                return self.returncode
+
+            def communicate(self):
+                return b"", b""
+
+        process = RunningProcess()
+
+        def fake_popen(args, **kwargs):
+            captured_paths.extend(
+                [
+                    args[args.index("-K") + 1],
+                    args[args.index("--output") + 1],
+                ]
+            )
+            stopped["value"] = True
+            return process
+
+        with mock.patch(
+            "app.ai._find_system_curl",
+            return_value=r"C:\Windows\System32\curl.exe",
+        ), mock.patch(
+            "app.ai.subprocess.Popen",
+            side_effect=fake_popen,
+        ):
+            with self.assertRaises(CancelledError):
+                ai._download_cmhub_image_with_curl(
+                    "https://cdn.example.com/generated.png",
+                    connect_timeout=3,
+                    read_timeout=30,
+                    max_bytes=ai.CMHUB_IMAGE_MAX_BYTES,
+                    should_stop=lambda: stopped["value"],
+                )
+
+        self.assertTrue(process.terminated)
+        self.assertFalse(process.killed)
+        self.assertTrue(captured_paths)
+        self.assertTrue(all(not os.path.exists(path) for path in captured_paths))
 
     def test_cmhub_upstream_error_retries_and_keeps_metadata(self):
         with self.make_temp_dir() as temp_dir:

tests/test_image_studio_generation.py

@@ -1,7 +1,10 @@
 import io
 import os
 import sys
+import threading
+import time
 import unittest
+from concurrent.futures import CancelledError
 from unittest import mock
 
 sys.path.insert(0, os.path.dirname(__file__))
@@ -193,7 +196,14 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
             stopped = {"value": False}
             saved_paths = []
 
-            def fake_download(request_result, out_path, config, on_event, job_id):
+            def fake_download(
+                request_result,
+                out_path,
+                config,
+                on_event,
+                job_id,
+                should_stop=None,
+            ):
                 os.makedirs(os.path.dirname(out_path), exist_ok=True)
                 with open(out_path, "wb") as fh:
                     fh.write(self._png_bytes())
@@ -235,6 +245,138 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
 
         self.assert_removed(temp_dir)
 
+    def test_download_cancel_maps_to_cancelled_and_keeps_resume_action(self):
+        with self.make_temp_dir() as temp_dir:
+            cfg, project, source = self._project_source(temp_dir)
+            stopped = {"value": False}
+
+            def fake_download(*args, **kwargs):
+                self.assertIsNotNone(kwargs.get("should_stop"))
+                stopped["value"] = True
+                raise CancelledError()
+
+            with mock.patch(
+                "app.image_studio_generation._runtime",
+                return_value=self._runtime(),
+            ), mock.patch(
+                "app.image_studio_generation.ai._cmhub_call_with_retry",
+                return_value={"task_id": "cmhub-cancel-download", "status": "queued"},
+            ), mock.patch(
+                "app.image_studio_generation.ai._cmhub_call_once",
+                return_value={
+                    "task_id": "cmhub-cancel-download",
+                    "status": "succeeded",
+                    "result": {
+                        "image_url": "https://cdn.example.com/cancel-download.png"
+                    },
+                },
+            ), mock.patch(
+                "app.image_studio_generation.ai._download_cmhub_image_with_retry",
+                side_effect=fake_download,
+            ):
+                summary = image_studio_generation.generate_image_jobs(
+                    project.id,
+                    source.id,
+                    "下载停止测试",
+                    1,
+                    config=cfg,
+                    path=cfg["db_path"],
+                    should_stop=lambda: stopped["value"],
+                )
+
+            self.assertEqual(1, summary["cancelled"])
+            job = summary["jobs"][0]["job"]
+            stored = image_studio.get_job(job.id, path=cfg["db_path"])
+            self.assertEqual("cancelled", stored.status)
+            self.assertEqual(image_studio.JOB_RECOVERY_RESUME, stored.recovery_action)
+            self.assertEqual(
+                [],
+                image_studio.list_assets(
+                    project.id,
+                    kind="generated_main",
+                    path=cfg["db_path"],
+                ),
+            )
+
+        self.assert_removed(temp_dir)
+
+    def test_run_jobs_cancels_queued_future_before_running_slot_is_released(self):
+        with self.make_temp_dir() as temp_dir:
+            cfg, project, source = self._project_source(temp_dir)
+            cfg["ai"]["image_concurrency"] = 1
+            jobs = [
+                image_studio.create_job(
+                    project.id,
+                    source_asset_id=source.id,
+                    job_type="main",
+                    prompt="排队停止测试",
+                    path=cfg["db_path"],
+                )
+                for _ in range(2)
+            ]
+            started_ids = []
+            first_started = threading.Event()
+            release_first = threading.Event()
+            stop_requested = threading.Event()
+            result_holder = {}
+
+            def fake_run_one(*args):
+                job_id = int(args[0])
+                started_ids.append(job_id)
+                if len(started_ids) == 1:
+                    first_started.set()
+                    release_first.wait(timeout=3)
+                updated = image_studio.update_job_status(
+                    job_id,
+                    "cancelled",
+                    error="用户停止",
+                    recovery_action=image_studio.JOB_RECOVERY_REGENERATE,
+                    path=cfg["db_path"],
+                )
+                return {
+                    "job": updated,
+                    "status": "cancelled",
+                    "error": "用户停止",
+                }
+
+            def run():
+                result_holder["summary"] = image_studio_generation.run_jobs(
+                    jobs,
+                    config=cfg,
+                    path=cfg["db_path"],
+                    should_stop=stop_requested.is_set,
+                )
+
+            with mock.patch(
+                "app.image_studio_generation._runtime",
+                return_value=self._runtime(),
+            ), mock.patch(
+                "app.image_studio_generation._run_one_job_with_global_slot",
+                side_effect=fake_run_one,
+            ):
+                thread = threading.Thread(target=run)
+                thread.start()
+                self.assertTrue(first_started.wait(timeout=2))
+                stop_requested.set()
+                time.sleep(0.35)
+                release_first.set()
+                thread.join(timeout=3)
+
+            self.assertFalse(thread.is_alive())
+            self.assertEqual([jobs[0].id], started_ids)
+            summary = result_holder["summary"]
+            self.assertEqual(2, summary["cancelled"])
+            self.assertEqual(2, len(summary["jobs"]))
+            self.assertEqual(
+                ["cancelled", "cancelled"],
+                [
+                    image_studio.get_job(job.id, path=cfg["db_path"]).status
+                    for job in jobs
+                ],
+            )
+
+        self.assert_removed(temp_dir)
+
     def test_resume_existing_job_polls_without_new_submit(self):
         with self.make_temp_dir() as temp_dir:
             cfg, project, source = self._project_source(temp_dir)

tests/test_product_suite_gui.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import time
 import unittest
 from unittest import mock
 
@@ -823,6 +824,251 @@ class ProductSuiteGuiTests(TempDirMixin, unittest.TestCase):
         self.app.processEvents()
         self.assertGreater(edit.height(), wide_height)
 
+    def test_generation_terminal_watchdog_finalizes_once_and_restores_button(self):
+        with self.make_temp_dir() as temp_dir:
+            config = self._config(temp_dir)
+            project, assets = self._create_project_with_assets(temp_dir, config, 1)
+            jobs = [
+                image_studio.create_job(
+                    project.id,
+                    source_asset_id=assets[0].id,
+                    job_type="白底图",
+                    prompt="终态看门狗测试",
+                    path=config["db_path"],
+                )
+                for _ in range(2)
+            ]
+            for job in jobs:
+                image_studio.update_job_status(
+                    job.id,
+                    "succeeded",
+                    path=config["db_path"],
+                )
+
+            tab = ProductSuiteTab(config=config, db_path=config["db_path"])
+            self.addCleanup(tab.close)
+            state = tab._displayed_state
+            state.account_alias = "alias-a"
+            state.item_id = project.item_id
+            state.project_id = project.id
+            state.project_binding_state = project.binding_state
+            state.worker = mock.Mock()
+            state.thread = mock.Mock()
+            state.generation_run_token = "watchdog-run"
+            state.current_job_ids = [job.id for job in jobs]
+            state.total = len(jobs)
+            state.started_at = time.monotonic()
+            tab._generation_run_states["watchdog-run"] = state.key
+            tab._load_state(state)
+            messages = []
+            tab._message = lambda title, message, **kwargs: messages.append(
+                (title, message)
+            )
+
+            tab._check_generation_watchdogs()
+            self.assertIsNotNone(state.worker)
+            tab._check_generation_watchdogs()
+
+            self.assertIsNone(state.worker)
+            self.assertIsNone(state.thread)
+            self.assertEqual("", state.generation_run_token)
+            self.assertTrue(tab.generate_button.text().startswith("生成套图"))
+            self.assertEqual(1, len(messages))
+            self.assertEqual("商品套图生成完成", messages[0][0])
+            self.assertFalse(
+                tab._finalize_generation(
+                    state,
+                    "watchdog-run",
+                    {"total": 2, "success": 2},
+                    source="worker",
+                )
+            )
+            self.assertEqual(1, len(messages))
+
+        self.assert_removed(temp_dir)
+
+    def test_generation_real_qthread_completion_restores_gui_state(self):
+        with self.make_temp_dir() as temp_dir:
+            config = self._config(temp_dir)
+            project, assets = self._create_project_with_assets(temp_dir, config, 1)
+            tab = ProductSuiteTab(config=config, db_path=config["db_path"])
+            self.addCleanup(tab.close)
+            state = tab._displayed_state
+            state.account_alias = "alias-a"
+            state.item_id = project.item_id
+            state.project_id = project.id
+            state.project_binding_state = project.binding_state
+            state.prompt = "真实线程完成测试"
+            tab._load_state(state)
+            messages = []
+            tab._message = lambda title, message, **kwargs: messages.append(
+                (title, message)
+            )
+
+            def fake_run_jobs(jobs, **kwargs):
+                job_list = list(jobs)
+                for job in job_list:
+                    image_studio.update_job_status(
+                        job.id,
+                        "succeeded",
+                        path=config["db_path"],
+                    )
+                return {
+                    "total": len(job_list),
+                    "success": len(job_list),
+                    "failed": 0,
+                    "cancelled": 0,
+                    "jobs": [],
+                }
+
+            with mock.patch(
+                "app.gui.workers.image_studio_generation.run_jobs",
+                side_effect=fake_run_jobs,
+            ):
+                self.assertTrue(
+                    tab.start_generation(
+                        state,
+                        specs=[
+                            {
+                                "source_asset_id": assets[0].id,
+                                "job_type": "白底图",
+                                "prompt": "真实线程完成测试",
+                            }
+                        ],
+                    )
+                )
+                deadline = time.monotonic() + 3
+                while state.worker is not None and time.monotonic() < deadline:
+                    QTest.qWait(20)
+                    self.app.processEvents()
+
+            self.assertIsNone(state.worker)
+            self.assertIsNone(state.thread)
+            self.assertTrue(tab.generate_button.text().startswith("生成套图"))
+            self.assertEqual(1, len(messages))
+            self.assertEqual("商品套图生成完成", messages[0][0])
+
+        self.assert_removed(temp_dir)
+
+    def test_generation_thread_finished_reconciles_nonterminal_job(self):
+        with self.make_temp_dir() as temp_dir:
+            config = self._config(temp_dir)
+            project, assets = self._create_project_with_assets(temp_dir, config, 1)
+            job = image_studio.create_job(
+                project.id,
+                source_asset_id=assets[0].id,
+                job_type="场景图",
+                prompt="线程结束兜底测试",
+                path=config["db_path"],
+            )
+            tab = ProductSuiteTab(config=config, db_path=config["db_path"])
+            self.addCleanup(tab.close)
+            state = tab._displayed_state
+            state.account_alias = "alias-a"
+            state.item_id = project.item_id
+            state.project_id = project.id
+            state.project_binding_state = project.binding_state
+            state.worker = mock.Mock()
+            state.thread = mock.Mock()
+            state.generation_run_token = "thread-fallback"
+            state.current_job_ids = [job.id]
+            state.total = 1
+            state.started_at = time.monotonic()
+            tab._generation_run_states["thread-fallback"] = state.key
+            tab._load_state(state)
+            messages = []
+            tab._message = lambda title, message, **kwargs: messages.append(
+                (title, message)
+            )
+
+            tab._handle_generation_thread_finished("thread-fallback")
+
+            stored = image_studio.get_job(job.id, path=config["db_path"])
+            self.assertEqual("cancelled", stored.status)
+            self.assertEqual(
+                image_studio.JOB_RECOVERY_REGENERATE,
+                stored.recovery_action,
+            )
+            self.assertIsNone(state.worker)
+            self.assertEqual("商品套图生成未完整结束", messages[0][0])
+            self.assertIn("稍后继续查询", messages[0][1])
+
+        self.assert_removed(temp_dir)
+
+    def test_generation_old_run_token_and_repeated_stop_are_ignored(self):
+        with self.make_temp_dir() as temp_dir:
+            config = self._config(temp_dir)
+            accounts.create_account("主店", "alias-a", debug_port=9222, config=config)
+            statuses = []
+            tab = ProductSuiteTab(
+                config=config,
+                db_path=config["db_path"],
+                status_callback=lambda message, level=None: statuses.append(
+                    (message, level)
+                ),
+            )
+            self.addCleanup(tab.close)
+            state = tab._displayed_state
+            state.worker = mock.Mock()
+            state.thread = mock.Mock()
+            state.generation_run_token = "current-run"
+            state.generation_stop_requested = True
+            state.total = 2
+            tab._generation_run_states["current-run"] = state.key
+            original_worker = state.worker
+
+            self.assertFalse(
+                tab._finalize_generation(
+                    state,
+                    "old-run",
+                    {"total": 2, "success": 2},
+                    source="worker",
+                )
+            )
+            self.assertIs(original_worker, state.worker)
+
+            confirm = mock.Mock(return_value=True)
+            with mock.patch.object(tab, "_confirm", confirm):
+                tab.toggle_generation()
+            confirm.assert_not_called()
+            original_worker.cancel.assert_not_called()
+            self.assertEqual(("正在停止当前套图任务", "warning"), statuses[-1])
+
+        self.assert_removed(temp_dir)
+
+    def test_generation_immediate_stop_before_job_creation_finishes_cleanly(self):
+        with self.make_temp_dir() as temp_dir:
+            config = self._config(temp_dir)
+            accounts.create_account("主店", "alias-a", debug_port=9222, config=config)
+            tab = ProductSuiteTab(config=config, db_path=config["db_path"])
+            self.addCleanup(tab.close)
+            state = tab._displayed_state
+            state.worker = mock.Mock()
+            state.thread = mock.Mock()
+            state.generation_run_token = "immediate-stop"
+            state.generation_stop_requested = True
+            state.total = 3
+            state.started_at = time.monotonic()
+            tab._generation_run_states["immediate-stop"] = state.key
+            messages = []
+            tab._message = lambda title, message, **kwargs: messages.append(
+                (title, message)
+            )
+
+            tab._on_generation_finished_signal(
+                {
+                    "run_token": "immediate-stop",
+                    "cancelled": True,
+                }
+            )
+
+            self.assertIsNone(state.worker)
+            self.assertEqual(3, state.done)
+            self.assertEqual("商品套图生成已停止", messages[0][0])
+            self.assertIn("停止3张", messages[0][1])
+
+        self.assert_removed(temp_dir)
+
     def test_project_settings_and_result_history_use_existing_backend(self):
         with self.make_temp_dir() as temp_dir:
             config = self._config(temp_dir)

tests/test_product_suite_workers.py

@@ -48,10 +48,13 @@ class ProductSuiteWorkerTests(TempDirMixin, unittest.TestCase):
                         "prompt": "通勤场景图",
                     },
                 ],
+                run_token="run-123",
                 aspect_ratio="3:4",
                 db_path=db_path,
                 config={"db_path": db_path},
             )
+            progress_events = []
+            worker.progress.connect(progress_events.append)
 
             with mock.patch(
                 "app.gui.workers.image_studio_generation.run_jobs",
@@ -71,6 +74,12 @@ class ProductSuiteWorkerTests(TempDirMixin, unittest.TestCase):
             self.assertEqual([job.id for job in reversed(jobs)], result["job_ids"])
             self.assertEqual("3:4", run_jobs.call_args.kwargs["aspect_ratio"])
             self.assertEqual(db_path, run_jobs.call_args.kwargs["path"])
+            self.assertTrue(progress_events)
+            self.assertTrue(
+                all(event.get("run_token") == "run-123" for event in progress_events)
+            )
+            self.assertEqual("run-123", result["run_token"])
+            self.assertEqual(0, result["cancelled_count"])
 
         self.assert_removed(temp_dir)