From b68e22c0944abd36308147550487098064d34a23 Mon Sep 17 00:00:00 2001 From: chengma Date: Thu, 16 Jul 2026 17:12:41 +0800 Subject: [PATCH] fix(product-suite): finalize and cancel generation --- app/ai.py | 124 ++++++-- app/gui/tabs/product_suite.py | 420 ++++++++++++++++++++++++-- app/gui/workers.py | 16 +- app/image_studio_generation.py | 75 ++++- docs/04-architecture.md | 1 + docs/routes.md | 2 +- docs/tasks/T-639.md | 7 +- tests/test_ai.py | 135 ++++++++- tests/test_image_studio_generation.py | 144 ++++++++- tests/test_product_suite_gui.py | 246 +++++++++++++++ tests/test_product_suite_workers.py | 9 + 11 files changed, 1094 insertions(+), 85 deletions(-) diff --git a/app/ai.py b/app/ai.py index 81af8ce..ce2077e 100644 --- a/app/ai.py +++ b/app/ai.py @@ -1447,6 +1447,7 @@ def _download_cmhub_image_with_retry( use_system_proxy=False, download_with_curl="false", on_step=None, + should_stop=None, attempts=CMHUB_IMAGE_DOWNLOAD_ATTEMPTS, slow_threshold=CMHUB_IMAGE_SLOW_DOWNLOAD_SECONDS, ): @@ -1454,6 +1455,7 @@ def _download_cmhub_image_with_retry( total_started = time.perf_counter() last_exc = None for index in range(total_attempts): + _raise_if_download_cancelled(should_stop) try: image_bytes = _download_cmhub_image( url, @@ -1461,7 +1463,9 @@ def _download_cmhub_image_with_retry( read_timeout=read_timeout, use_system_proxy=use_system_proxy, download_with_curl=download_with_curl, + should_stop=should_stop, ) + _raise_if_download_cancelled(should_stop) elapsed = time.perf_counter() - total_started if elapsed >= float(slow_threshold or 0): _notify_step_event( @@ -1473,6 +1477,8 @@ def _download_cmhub_image_with_retry( level="warning", ) return image_bytes, elapsed + except CancelledError: + raise except Exception as exc: last_exc = exc if index + 1 >= total_attempts or not _cmhub_download_retryable(exc): @@ -1486,7 +1492,7 @@ def _download_cmhub_image_with_retry( attempt=index + 1, attempts=total_attempts, ) - time.sleep(min(2.0, 0.5 * (index + 1))) + _sleep_download_retry(min(2.0, 0.5 * (index + 1)), should_stop) if total_attempts > 1 and _cmhub_download_retryable(last_exc): raise AIError( "下载 cmhub 图片失败（已尝试 %s 次）: %s" @@ -1739,7 +1745,9 @@ def _download_cmhub_image( max_bytes=CMHUB_IMAGE_MAX_BYTES, use_system_proxy=False, download_with_curl="false", + should_stop=None, ): + _raise_if_download_cancelled(should_stop) _assert_public_http_url(url) if _should_use_curl_for_cmhub_download(download_with_curl): try: @@ -1749,7 +1757,10 @@ def _download_cmhub_image( read_timeout=read_timeout, max_bytes=max_bytes, use_system_proxy=use_system_proxy, + should_stop=should_stop, ) + except CancelledError: + raise except AIError: pass return _download_cmhub_image_with_requests( @@ -1757,32 +1768,53 @@ def _download_cmhub_image( connect_timeout=connect_timeout, read_timeout=read_timeout, max_bytes=max_bytes, + should_stop=should_stop, ) -def _download_cmhub_image_with_requests(url, connect_timeout, read_timeout, max_bytes): +def _download_cmhub_image_with_requests( + url, + connect_timeout, + read_timeout, + max_bytes, + should_stop=None, +): + response = None try: + _raise_if_download_cancelled(should_stop) response = _cmhub_session().get( url, stream=True, timeout=(max(1, int(connect_timeout)), max(1, int(read_timeout))), ) + _raise_if_download_cancelled(should_stop) + status = getattr(response, "status_code", 200) + if status >= 400: + raise AIError("下载 cmhub 图片失败: HTTP %s" % status) + chunks = [] + total = 0 + iterator = ( + response.iter_content(chunk_size=65536) + if hasattr(response, "iter_content") + else [response.content] + ) + for chunk in iterator: + _raise_if_download_cancelled(should_stop) + if not chunk: + continue + total += len(chunk) + if total > max_bytes: + raise AIError("下载 cmhub 图片失败: 图片超过大小上限") + chunks.append(chunk) + _raise_if_download_cancelled(should_stop) + return b"".join(chunks) + except CancelledError: + raise except requests.exceptions.RequestException as exc: raise AIError("下载 cmhub 图片失败: %s" % exc) from exc - status = getattr(response, "status_code", 200) - if status >= 400: - raise AIError("下载 cmhub 图片失败: HTTP %s" % status) - chunks = [] - total = 0 - iterator = response.iter_content(chunk_size=65536) if hasattr(response, "iter_content") else [response.content] - for chunk in iterator: - if not chunk: - continue - total += len(chunk) - if total > max_bytes: - raise AIError("下载 cmhub 图片失败: 图片超过大小上限") - chunks.append(chunk) - return b"".join(chunks) + finally: + if response is not None and hasattr(response, "close"): + response.close() def _should_use_curl_for_cmhub_download(mode): @@ -1846,6 +1878,7 @@ def _download_cmhub_image_with_curl( read_timeout, max_bytes, use_system_proxy=False, + should_stop=None, ): curl_path = _find_system_curl() if not curl_path: @@ -1881,20 +1914,36 @@ def _download_cmhub_image_with_curl( ] if not bool(use_system_proxy): args.extend(["--noproxy", "*"]) + process = None try: - completed = subprocess.run( + process = subprocess.Popen( args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, - timeout=max(2, int(connect_timeout) + int(read_timeout) + 10), - check=False, shell=False, **_subprocess_hidden_window_kwargs(), ) - except (OSError, subprocess.TimeoutExpired) as exc: + deadline = time.monotonic() + max( + 2, + int(connect_timeout) + int(read_timeout) + 10, + ) + while process.poll() is None: + try: + _raise_if_download_cancelled(should_stop) + except CancelledError: + _stop_download_process(process) + raise + if time.monotonic() >= deadline: + _stop_download_process(process) + raise AIError("下载 cmhub 图片失败: curl 执行超时") + time.sleep(0.1) + process.communicate() + except CancelledError: + raise + except OSError as exc: raise AIError("下载 cmhub 图片失败: curl 执行失败") from exc - if completed.returncode != 0: - raise AIError("下载 cmhub 图片失败: curl 退出码 %s" % completed.returncode) + if process.returncode != 0: + raise AIError("下载 cmhub 图片失败: curl 退出码 %s" % process.returncode) size = os.path.getsize(temp_output_path) if size > max_bytes: raise AIError("下载 cmhub 图片失败: 图片超过大小上限") @@ -1909,6 +1958,37 @@ def _download_cmhub_image_with_curl( pass +def _stop_download_process(process): + if process is None or process.poll() is not None: + return + try: + process.terminate() + process.wait(timeout=2) + except (OSError, subprocess.TimeoutExpired): + try: + process.kill() + process.wait(timeout=2) + except (OSError, subprocess.TimeoutExpired): + pass + + +def _raise_if_download_cancelled(should_stop): + try: + stopped = bool(should_stop and should_stop()) + except Exception: + stopped = False + if stopped: + raise CancelledError() + + +def _sleep_download_retry(delay_seconds, should_stop): + deadline = time.monotonic() + max(0.0, float(delay_seconds or 0.0)) + while time.monotonic() < deadline: + _raise_if_download_cancelled(should_stop) + time.sleep(min(0.1, max(0.0, deadline - time.monotonic()))) + _raise_if_download_cancelled(should_stop) + + def _curl_config_quote(value): text = str(value or "") return '"' + text.replace("\\", "\\\\").replace('"', '\\"') + '"' diff --git a/app/gui/tabs/product_suite.py b/app/gui/tabs/product_suite.py index 8a65a93..cfdaed6 100644 --- a/app/gui/tabs/product_suite.py +++ b/app/gui/tabs/product_suite.py @@ -5,9 +5,20 @@ from __future__ import annotations import os import re import time +import uuid from dataclasses import dataclass, field -from PySide6.QtCore import QByteArray, QBuffer, QIODevice, QRect, QSize, Qt, QTimer, Signal +from PySide6.QtCore import ( + QByteArray, + QBuffer, + QIODevice, + QRect, + QSize, + Qt, + QTimer, + Signal, + Slot, +) from PySide6.QtGui import QColor, QIcon, QImage, QImageReader, QKeySequence, QPainter, QPixmap from PySide6.QtWidgets import ( QApplication, @@ -688,6 +699,9 @@ class SuiteTaskState: show_history: bool = False worker: object = None thread: object = None + generation_run_token: str = "" + generation_stop_requested: bool = False + generation_terminal_streak: int = 0 pull_worker: object = None pull_thread: object = None import_worker: object = None @@ -739,6 +753,7 @@ class ProductSuiteTab(QWidget): self._next_serial = 1 self._displayed_state = None self._prompt_save_timers = {} + self._generation_run_states = {} self._original_list_context = None self._loading = False self._result_refresh_pending = False @@ -759,6 +774,10 @@ class ProductSuiteTab(QWidget): self.elapsed_timer.setInterval(1000) self.elapsed_timer.timeout.connect(self._refresh_elapsed) self.elapsed_timer.start() + self.generation_watchdog_timer = QTimer(self) + self.generation_watchdog_timer.setInterval(750) + self.generation_watchdog_timer.timeout.connect(self._check_generation_watchdogs) + self.generation_watchdog_timer.start() if self._prompt_template_init_error: self._status(self._prompt_template_init_error, "danger") @@ -1324,6 +1343,7 @@ class ProductSuiteTab(QWidget): destructive=True, ): return + state.generation_stop_requested = True state.worker.cancel() if state.ai_worker is not None: state.ai_worker.cancel() @@ -1347,6 +1367,8 @@ class ProductSuiteTab(QWidget): except Exception as exc: self._status("清理空临时草稿失败：%s" % _user_error(exc), "danger") self._release_prompt_save_timer(state) + if state.generation_run_token: + self._generation_run_states.pop(state.generation_run_token, None) self._retired_states.append(state) self._states.pop(state.key, None) self.task_tabs.removeTab(index) @@ -1884,6 +1906,10 @@ class ProductSuiteTab(QWidget): token = id(thread) _PRODUCT_SUITE_THREAD_REFS[token] = (thread, worker) thread.finished.connect(lambda token=token: _PRODUCT_SUITE_THREAD_REFS.pop(token, None)) + run_token = str(getattr(worker, "run_token", "") or "") + if run_token: + thread.setProperty("productSuiteRunToken", run_token) + thread.finished.connect(self._on_generation_thread_finished_signal) thread.start() return thread @@ -2491,14 +2517,22 @@ class ProductSuiteTab(QWidget): if state is None: return if state.generation_running(): + if state.generation_stop_requested: + self._status("正在停止当前套图任务", "warning") + return if self._confirm( "停止生成套图", "确认取消当前任务吗？已提交任务会在安全边界停止。", destructive=True, ): + state.generation_stop_requested = True state.worker.cancel() - self.generate_button.setText("正在停止...") - self.generate_button.setEnabled(False) + self._log_generation_lifecycle( + state, + state.generation_run_token, + "stop_requested", + ) + self._apply_running_state(state) self._status("已请求停止当前套图任务", "warning") return self.start_generation(state) @@ -2547,26 +2581,38 @@ class ProductSuiteTab(QWidget): ): return False self._persist_state(state) + run_token = uuid.uuid4().hex worker = ProductSuiteGenerateWorker( state.project_id, specs, + run_token=run_token, aspect_ratio=state.settings["ratio"], db_path=self.db_path, config=self.config, cmhub_config_path=self.cmhub_config_path, ) state.worker = worker + state.generation_run_token = run_token + state.generation_stop_requested = False + state.generation_terminal_streak = 0 state.done = 0 state.failed = 0 state.total = len(specs) state.started_at = time.monotonic() state.current_job_ids = [] state.show_history = False - worker.progress.connect(lambda payload, state=state: self._on_generation_progress(state, payload)) - worker.finished.connect(lambda result, state=state: self._on_generation_finished(state, result)) - worker.cancelled.connect(lambda result, state=state: self._on_generation_finished(state, result)) - worker.failed.connect(lambda row, error, state=state: self._on_generation_failed(state, error)) + self._generation_run_states[run_token] = state.key + worker.progress.connect(self._on_generation_progress_signal) + worker.finished.connect(self._on_generation_finished_signal) + worker.cancelled.connect(self._on_generation_finished_signal) + worker.failed.connect(self._on_generation_failed_signal) state.thread = self._start_thread(worker, "商品套图生成") + self._log_generation_lifecycle( + state, + run_token, + "started", + {"total": len(specs), "job_ids": 0}, + ) if state is self._displayed_state: self._loading = True try: @@ -2578,52 +2624,367 @@ class ProductSuiteTab(QWidget): self._status("商品套图生成已开始，共%d张；可切换到其他任务" % len(specs), "info") return True + def _generation_signal_token(self, payload=None): + token = str((payload or {}).get("run_token") or "") + sender = self.sender() + return token or str(getattr(sender, "run_token", "") or "") + + def _generation_state(self, run_token): + token = str(run_token or "") + state = self._states.get(self._generation_run_states.get(token)) + if state is None or state.generation_run_token != token: + return None + return state + + @Slot(dict) + def _on_generation_progress_signal(self, payload): + token = self._generation_signal_token(payload) + state = self._generation_state(token) + if state is None: + return + self._on_generation_progress(state, payload) + def _on_generation_progress(self, state, payload): - state.total = int(payload.get("total", state.total) or state.total) + if "total" in payload: + state.total = max(0, int(payload.get("total") or 0)) state.done = int(payload.get("done", state.done) or 0) state.failed = int(payload.get("failed", state.failed) or 0) job_ids = payload.get("job_ids") - if job_ids: + if job_ids is not None: state.current_job_ids = [int(job_id) for job_id in job_ids] + state.generation_terminal_streak = 0 if state is self._displayed_state: self._refresh_results(state) self._refresh_elapsed() - def _on_generation_failed(self, state, error): + @Slot(int, str) + def _on_generation_failed_signal(self, row, error): + token = self._generation_signal_token() + state = self._generation_state(token) + if state is None: + return + self._log_generation_lifecycle( + state, + token, + "worker_failed", + {"has_error": True}, + level="ERROR", + ) self._status("商品套图生成失败：%s" % _user_error(error), "danger") - def _on_generation_finished(self, state, result): + @Slot(dict) + def _on_generation_finished_signal(self, result): + token = self._generation_signal_token(result) + state = self._generation_state(token) + if state is None: + return + self._log_generation_lifecycle( + state, + token, + "worker_finished", + { + "ok": result.get("ok", True), + "success": result.get("success", 0), + "failed": result.get("failed", 0), + "cancelled": result.get( + "cancelled_count", + result.get("cancelled", 0), + ), + }, + ) + if result.get("ok") is False: + self._reconcile_generation_jobs(state, "生成线程异常结束") + if ( + state.generation_stop_requested + and not self._generation_job_ids(state) + and result.get("cancelled") is True + ): + result = dict(result) + result["total"] = state.total + result["cancelled_count"] = state.total + self._finalize_generation(state, token, result, source="worker") + + @Slot() + def _on_generation_thread_finished_signal(self): + sender = self.sender() + token = str( + sender.property("productSuiteRunToken") + if sender is not None + else "" + ) + self._handle_generation_thread_finished(token) + + def _handle_generation_thread_finished(self, run_token): + state = self._generation_state(run_token) + if state is None: + return + snapshot = self._generation_job_snapshot(state) + had_active = bool(snapshot["active"]) + self._log_generation_lifecycle( + state, + run_token, + "thread_finished_fallback", + snapshot, + level="WARNING" if snapshot["active"] else "INFO", + ) + if had_active: + self._reconcile_generation_jobs(state, "生成线程已结束") + snapshot = self._generation_job_snapshot(state) + result = self._generation_result_from_snapshot(snapshot) + if not snapshot["job_ids"] and state.generation_stop_requested: + result.update( + { + "total": state.total, + "cancelled_count": state.total, + } + ) + elif had_active: + result.update( + { + "ok": False, + "error": "生成线程已结束，部分任务可稍后继续查询", + } + ) + self._finalize_generation( + state, + run_token, + result, + source="thread_finished", + ) + + def _check_generation_watchdogs(self): + for state in list(self._states.values()): + token = state.generation_run_token + if not token or not state.generation_running(): + continue + snapshot = self._generation_job_snapshot(state) + if snapshot["all_terminal"]: + state.generation_terminal_streak += 1 + else: + state.generation_terminal_streak = 0 + if state.generation_terminal_streak < 2: + continue + self._log_generation_lifecycle( + state, + token, + "terminal_watchdog_finalize", + snapshot, + level="WARNING", + ) + self._finalize_generation( + state, + token, + self._generation_result_from_snapshot(snapshot), + source="terminal_watchdog", + ) + + def _generation_job_ids(self, state): + job_ids = list(state.current_job_ids) + if not job_ids and state.worker is not None: + worker_job_ids = getattr(state.worker, "job_ids", []) + if not isinstance(worker_job_ids, (list, tuple, set)): + worker_job_ids = [] + job_ids = [ + int(job_id) + for job_id in list(worker_job_ids or []) + ] + if job_ids: + state.current_job_ids = job_ids + return job_ids + + def _generation_job_snapshot(self, state): + job_ids = self._generation_job_ids(state) + counts = { + "success": 0, + "failed": 0, + "cancelled": 0, + "active": 0, + "job_ids": len(job_ids), + "all_terminal": False, + } + if not job_ids: + counts["active"] = max(0, int(state.total or 0)) + return counts + for job_id in job_ids: + try: + job = image_studio.get_job(job_id, path=self.db_path) + except Exception: + job = None + status = str(getattr(job, "status", "") or "") + if status == "succeeded": + counts["success"] += 1 + elif status in {"failed", "expired"}: + counts["failed"] += 1 + elif status == "cancelled": + counts["cancelled"] += 1 + else: + counts["active"] += 1 + counts["all_terminal"] = ( + len(job_ids) == int(state.total or 0) + and counts["active"] == 0 + ) + return counts + + def _generation_result_from_snapshot(self, snapshot): + return { + "total": int(snapshot.get("job_ids", 0) or 0), + "success": int(snapshot.get("success", 0) or 0), + "failed": int(snapshot.get("failed", 0) or 0), + "cancelled_count": int(snapshot.get("cancelled", 0) or 0), + } + + def _reconcile_generation_jobs(self, state, reason): + for job_id in self._generation_job_ids(state): + try: + job = image_studio.get_job(job_id, path=self.db_path) + if job is None or job.status in { + "succeeded", + "failed", + "expired", + "cancelled", + }: + continue + recovery = ( + image_studio.JOB_RECOVERY_RESUME + if job.task_id + else image_studio.JOB_RECOVERY_REGENERATE + ) + image_studio.update_job_status( + job.id, + "cancelled", + error="%s，任务可稍后继续处理" % reason, + recovery_action=recovery, + path=self.db_path, + ) + except Exception as exc: + self._status( + "商品套图任务状态收尾失败：%s" % _user_error(exc), + "danger", + ) + + def _finalize_generation(self, state, run_token, result, *, source): + if self._generation_state(run_token) is not state: + return False + snapshot = self._generation_job_snapshot(state) + if snapshot["job_ids"] and ( + snapshot["all_terminal"] + or source in {"thread_finished", "terminal_watchdog"} + ): + result = dict(result or {}) + result.update(self._generation_result_from_snapshot(snapshot)) + result = dict(result or {}) + stop_requested = state.generation_stop_requested + success = int(result.get("success", 0) or 0) + failed = int(result.get("failed", 0) or 0) + cancelled = int( + result.get("cancelled_count", result.get("cancelled", 0)) or 0 + ) + total = int(result.get("total", state.total) or state.total) + active = max(0, total - success - failed - cancelled) + elapsed = ( + int(max(0, time.monotonic() - state.started_at)) + if state.started_at + else 0 + ) + self._generation_run_states.pop(run_token, None) + state.generation_run_token = "" + state.generation_stop_requested = False + state.generation_terminal_streak = 0 state.worker = None state.thread = None - state.done = int(result.get("success", 0) or 0) + int(result.get("failed", 0) or 0) + int( - result.get("cancelled", 0) or 0 - ) - state.failed = int(result.get("failed", state.failed) or 0) - state.total = int(result.get("total", state.total) or state.total) - elapsed = int(max(0, time.monotonic() - state.started_at)) if state.started_at else 0 + state.done = success + failed + cancelled + state.failed = failed + state.total = total state.started_at = None if state is self._displayed_state: self._apply_running_state(state) self._refresh_results(state) self._refresh_elapsed() + self._log_generation_lifecycle( + state, + run_token, + "finalized", + { + "source": source, + "total": total, + "success": success, + "failed": failed, + "cancelled": cancelled, + "active": active, + "elapsed_seconds": elapsed, + }, + level="WARNING" if active or result.get("ok") is False else "INFO", + ) if state.key not in self._states: - return - if result.get("ok") is False: + return True + if active or result.get("ok") is False: + message = str( + result.get("error") + or "生成线程已结束，部分任务可稍后继续查询" + ) if state is self._displayed_state: - self._message("商品套图生成失败", _user_error(result.get("error"))) + self._message("商品套图生成未完整结束", _user_error(message)) else: - self._status("套图任务%d生成失败" % state.serial, "danger") - return - success = int(result.get("success", 0) or 0) - cancelled = int(result.get("cancelled", 0) or 0) + self._status( + "套图任务%d生成未完整结束" % state.serial, + "danger", + ) + return True + if stop_requested or cancelled: + if state is self._displayed_state: + self._message( + "商品套图生成已停止", + "本轮共%d张：成功%d张，失败%d张，停止%d张；" + "已提交任务可稍后继续查询；总用时%d秒。" + % (total, success, failed, cancelled, elapsed), + icon=QMessageBox.Information, + ) + self._status( + "商品套图生成已停止：成功%d张，失败%d张，停止%d张" + % (success, failed, cancelled), + "warning", + ) + return True if state is self._displayed_state: self._message( "商品套图生成完成", "本轮共%d张：成功%d张，失败%d张，停止%d张；总用时%d秒。" - % (state.total, success, state.failed, cancelled, elapsed), + % (total, success, failed, cancelled, elapsed), icon=QMessageBox.Information, ) - self._status("商品套图生成完成：成功%d张，失败%d张" % (success, state.failed), "success") + self._status( + "商品套图生成完成：成功%d张，失败%d张" % (success, failed), + "success", + ) + return True + + def _log_generation_lifecycle( + self, + state, + run_token, + event, + payload=None, + *, + level="INFO", + ): + data = { + "run_token": str(run_token or "")[:8], + "project_id": getattr(state, "project_id", None), + "event": str(event or ""), + } + data.update(dict(payload or {})) + try: + diagnostics.write_diagnostic_log( + "商品套图生成生命周期", + level=level, + step="product_suite_generation", + task_id=getattr(state, "project_id", None), + item_id=getattr(state, "item_id", None), + payload=data, + log_dir=appconfig.diagnostic_log_dir(self.config), + ) + except Exception: + pass def _apply_running_state(self, state): generation_running = state.generation_running() @@ -2648,7 +3009,11 @@ class ProductSuiteTab(QWidget): row.set_controls_enabled(not generation_running) self.generate_button.setEnabled(True) if generation_running: - self.generate_button.setText("停止生成") + self.generate_button.setText( + "正在停止..." + if state.generation_stop_requested + else "停止生成" + ) self.generate_button.setStyleSheet( "QPushButton { background: #cf222e; color: white; border-color: #a40e26; font-weight: 600; }" "QPushButton:hover { background: #a40e26; }" @@ -2909,4 +3274,5 @@ class ProductSuiteTab(QWidget): worker.cancel() for state in list(self._states.values()) + list(self._retired_states): self._release_prompt_save_timer(state) + self._generation_run_states.clear() super().closeEvent(event) diff --git a/app/gui/workers.py b/app/gui/workers.py index f9c3bf7..2ce1d5d 100644 --- a/app/gui/workers.py +++ b/app/gui/workers.py @@ -303,6 +303,7 @@ class ProductSuiteGenerateWorker(BaseWorker): project_id, job_specs, *, + run_token="", aspect_ratio="1:1", db_path=None, config=None, @@ -311,10 +312,12 @@ class ProductSuiteGenerateWorker(BaseWorker): super().__init__() self.project_id = int(project_id) self.job_specs = [dict(spec) for spec in (job_specs or [])] + self.run_token = str(run_token or "") self.aspect_ratio = str(aspect_ratio or "1:1") self.db_path = db_path self.config = config self.cmhub_config_path = cmhub_config_path + self.job_ids = [] self._done = 0 self._failed = 0 self._lock = threading.Lock() @@ -325,8 +328,6 @@ class ProductSuiteGenerateWorker(BaseWorker): raise ValueError("商品套图生成任务不能为空") jobs = [] for spec in self.job_specs: - if self.should_cancel(): - break jobs.append( image_studio.create_job( self.project_id, @@ -338,12 +339,14 @@ class ProductSuiteGenerateWorker(BaseWorker): path=self.db_path, ) ) + self.job_ids = [job.id for job in jobs] self.progress.emit( { + "run_token": self.run_token, "total": len(jobs), "done": 0, "failed": 0, - "job_ids": [job.id for job in jobs], + "job_ids": list(self.job_ids), } ) @@ -356,10 +359,11 @@ class ProductSuiteGenerateWorker(BaseWorker): if event.get("result") != "success": self._failed += 1 progress = { + "run_token": self.run_token, "total": len(jobs), "done": self._done, "failed": self._failed, - "job_ids": [job.id for job in jobs], + "job_ids": list(self.job_ids), } self.progress.emit(progress) @@ -373,7 +377,9 @@ class ProductSuiteGenerateWorker(BaseWorker): on_event=on_event, ) summary["project_id"] = self.project_id - summary["job_ids"] = [job.id for job in jobs] + summary["job_ids"] = list(self.job_ids) + summary["cancelled_count"] = int(summary.get("cancelled", 0) or 0) + summary["run_token"] = self.run_token return summary diff --git a/app/image_studio_generation.py b/app/image_studio_generation.py index 6bd5623..2dae04c 100644 --- a/app/image_studio_generation.py +++ b/app/image_studio_generation.py @@ -6,7 +6,7 @@ import os import threading import time import urllib.parse -from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait +from concurrent.futures import CancelledError, FIRST_COMPLETED, ThreadPoolExecutor, wait from . import ai, appconfig, image_studio from .version import APP_VERSION @@ -185,7 +185,11 @@ def run_jobs( for job in job_list } while futures: - done, _ = wait(set(futures), return_when=FIRST_COMPLETED) + done, _ = wait( + set(futures), + timeout=0.2, + return_when=FIRST_COMPLETED, + ) for future in done: futures.pop(future) try: @@ -252,7 +256,14 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho request_result = _poll_job(job.id, request_result["task_id"], runtime, request_result, db_path, should_stop, on_event) _raise_if_stopped(should_stop) out_path = _output_path(project, job, image_root) - saved_path = _download_and_save_job_image(request_result, out_path, config, on_event, job.id) + saved_path = _download_and_save_job_image( + request_result, + out_path, + config, + on_event, + job.id, + should_stop=should_stop, + ) try: _raise_if_stopped(should_stop) except ImageStudioGenerationError: @@ -281,17 +292,27 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho _notify(on_event, {"job_id": job.id, "step": "job_done", "result": "success"}) return {"job": updated, "asset": asset, "status": "succeeded"} except Exception as exc: - status = "cancelled" if "停止" in str(exc) else "failed" + cancelled = isinstance(exc, CancelledError) or "停止" in str(exc) + status = "cancelled" if cancelled else "failed" + error = "用户已停止，已提交任务可稍后继续查询" if cancelled else str(exc) current_job = image_studio.get_job(job.id, path=db_path) updated = image_studio.update_job_status( job.id, status, - error=str(exc), + error=error, recovery_action=_recovery_action_for_job(current_job), path=db_path, ) - _notify(on_event, {"job_id": job.id, "step": "job_done", "result": status, "detail": str(exc)}) - return {"job": updated, "status": status, "error": str(exc)} + _notify( + on_event, + { + "job_id": job.id, + "step": "job_done", + "result": status, + "detail": error, + }, + ) + return {"job": updated, "status": status, "error": error} def _submit_or_resume_job( @@ -395,6 +416,7 @@ def _poll_job(job_id, task_id, runtime, request_result, db_path, should_stop, on read_timeout=ai.CMHUB_IMAGE_POLL_READ_TIMEOUT_SECONDS, headers_extra={"X-Client-Version": str(APP_VERSION)}, ) + _raise_if_stopped(should_stop) status = str(data.get("status") or "").strip().lower() if status in {"queued", "running"}: _notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": status, "task_id": task_id}) @@ -424,21 +446,44 @@ def _poll_job(job_id, task_id, runtime, request_result, db_path, should_stop, on raise ImageStudioGenerationError("cmhub 生图任务状态返回格式错误") -def _download_and_save_job_image(request_result, out_path, config, on_event, job_id): +def _download_and_save_job_image( + request_result, + out_path, + config, + on_event, + job_id, + should_stop=None, +): + _raise_if_stopped(should_stop) _notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "start"}) - image_bytes, _ = ai._download_cmhub_image_with_retry( - request_result["image_url"], - connect_timeout=request_result["connect_timeout"], - read_timeout=request_result["read_timeout"], - use_system_proxy=request_result.get("use_system_proxy", False), - download_with_curl=request_result.get("download_with_curl", "auto"), - ) + try: + image_bytes, _ = ai._download_cmhub_image_with_retry( + request_result["image_url"], + connect_timeout=request_result["connect_timeout"], + read_timeout=request_result["read_timeout"], + use_system_proxy=request_result.get("use_system_proxy", False), + download_with_curl=request_result.get("download_with_curl", "auto"), + should_stop=should_stop, + ) + except CancelledError as exc: + raise ImageStudioGenerationError( + "用户已停止，已提交任务可稍后继续查询" + ) from exc + _raise_if_stopped(should_stop) saved_path = ai._save_jpeg( image_bytes, out_path, request_result.get("resolution") or appconfig.ai_config(config).get("resolution", "1k"), request_result.get("quality") or appconfig.ai_config(config).get("jpg_quality", 90), ) + try: + _raise_if_stopped(should_stop) + except ImageStudioGenerationError: + try: + if os.path.isfile(saved_path): + os.remove(saved_path) + finally: + raise _notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "success"}) return saved_path diff --git a/docs/04-architecture.md b/docs/04-architecture.md index 5f65f7d..96da36f 100644 --- a/docs/04-architecture.md +++ b/docs/04-architecture.md @@ -437,6 +437,7 @@ data/images///__new. # AI 生成的新 - ⑥已选账号但未填写商品 ID 时允许导入、拖入或粘贴本地图片，首次有效导入才创建草稿；取消选择和全部导入失败不保留空草稿。草稿可管理本地图片、AI 帮写、生成套图、查看历史和打开结果目录，但在创建 worker、启动 Chrome 或执行 CDP 前禁止「拉取蝦皮主图」。输入合法数字商品 ID 后，经确认原地绑定同一个 `project_id`；资产、job、selection、提示词、套图设置和 `storage_key` 均保持不变。若同账号目标 ID（含软删除项目）已存在则拒绝覆盖或合并。 - 启动时恢复未软删除、至少含一条资产或生成任务的草稿为独立中文“临时草稿”标签，按最近更新时间排序。关闭非空草稿可选择保留、软删除或取消；软删除不物理删除图片目录。③「更新蝦皮」只处理正式任务，不接受临时草稿。 - 第六 Tab 的多个 `SuiteTaskState` 各自保留 generation/pull/import/AI/download worker 与线程引用；切换任务不取消任务。多个任务可并行，但 `image_studio_generation` 使用进程级 semaphore 保证所有套图任务合计最多5个 cmhub 在途 job。线程还在运行时关闭任务只请求协作式停止，模块级引用保留到 `QThread.finished`，不得提前销毁线程对象；下载前后均检查停止信号，停止后的临时文件不入资产库。 +- T-639 后每轮套图生成使用仅存在内存的 `run_token` 隔离迟到信号，progress/finished/cancelled/failed 通过主线程绑定槽统一处理；正常 worker 结果、`QThread.finished` 和本轮 job 连续两次全部终态看门狗共同进入幂等 finalize。GUI 只按本轮明确 `job_ids` 判断完成，不用历史图片数量；即使最终信号丢失也会恢复按钮，旧线程引用仍保留到真实结束。停止为协作式：调度循环约每200ms检查标记并取消未开始 future，提交/轮询在有界请求返回后停止；requests 在流式数据块边界取消，Windows curl 由隐藏窗口 `Popen` 有界 terminate/kill。已有 `task_id` 的停止任务保留 resume，不假设服务端任务被取消或点数退回。 提示词管理： diff --git a/docs/routes.md b/docs/routes.md index 685ae3d..89880b2 100644 --- a/docs/routes.md +++ b/docs/routes.md @@ -205,7 +205,7 @@ - 「拉取蝦皮主图」复用只读 CDP，读取 URL 后由最多2个下载 worker 后台落盘；不改标题/封面、不拖拽、不点击更新。拉取、下载期间其余界面和其他任务仍可操作。 - 套图只有一个图片类型，不再展示详情图、终选盘或模板 CRUD。默认分类为白底图1、场景图2、卖点图2；自定义分类名称非空、无空格、最多10字且不可重名。逐图主图开启后，白底图只生成一次，其余分类按每张有效原图展开。 - 平台、国家地区、语言和比例以四个带独立标签的同行下拉展示，选项只显示真实值；四项都写进每个 job 的完整提示词，比例还透传到 cmhub 生图请求，不是装饰字段。已有项目保存自己的完整设置；未绑定商品的新任务在重启后采用 `config.json` 的最近四项选择。生成仍走 `image_studio_generation.run_jobs()` 的 submit → poll → download 管线。 -- 生成按钮按当前总数显示并在运行时切换为停止。结果区显示本轮或历史 job；成功图可预览、复制路径、打开目录、重新生成、移入项目废纸篓并撤销，失败卡显示脱敏中文摘要与重试入口。 +- 生成按钮按当前总数显示并在运行时切换为「停止生成」；确认停止后显示「正在停止...」，重复点击不再弹确认框。每轮生成用独立运行标识隔离旧信号，本轮全部 job 终态或线程结束时都会统一恢复按钮；最终 worker 信号缺失时由数据库终态看门狗兜底，不要求用户重启。停止会取消未开始任务，已提交任务停止本地等待并保留后续继续查询语义；客户端不承诺取消服务端任务或退回点数。结果区显示本轮或历史 job；成功图可预览、复制路径、打开目录、重新生成、移入项目废纸篓并撤销，失败卡显示脱敏中文摘要与重试入口。 - AI帮写和生图按任务独立运行。AI帮写期间若用户改过卖点，返回后必须确认才覆盖；全部用户可见错误隐藏 URL/接口路径和敏感信息。 - ⑥只管理本地图片资产，不自动上传或修改蝦皮；③线上更新边界不受影响。旧 `ImageStudioTab` 留作代码兼容但不再作为主窗口入口。 diff --git a/docs/tasks/T-639.md b/docs/tasks/T-639.md index dae69e0..5c233b6 100644 --- a/docs/tasks/T-639.md +++ b/docs/tasks/T-639.md @@ -3,7 +3,7 @@ id: T-639 title: 商品套图生成完成状态复位与可响应停止 phase: 7 deps: [T-638] -status: TODO +status: DONE created: 2026-07-16 --- @@ -182,3 +182,8 @@ git diff --check - 不修改 CDP、蝦皮主图拉取、①导入采集、②AI生成、③更新蝦皮、④账号管理或⑤设置。 ## 执行记录 + +- 2026-07-16：每轮商品套图生成新增内存 `run_token`，worker 的 progress/summary 携带同一 token，GUI 改用主线程绑定槽处理信号。正常结果、`QThread.finished` 和每750ms检查的本轮 job 终态看门狗统一进入幂等 finalize；旧轮次迟到信号会被忽略，旧线程引用仍由模块级容器保留到真实结束。 +- 2026-07-16：统一 finalize 按本轮明确 `job_ids` 查询 SQLite，恢复按钮和编辑控件，并只显示一次中文完成/停止/异常汇总。线程异常结束时，已有 `task_id` 的非终态 job 记为 `cancelled + resume`，未提交任务记为 `cancelled + regenerate`；生命周期诊断日志只记录 token 短值、项目ID、数量、来源和用时，不记录提示词、URL或密钥。 +- 2026-07-16：`image_studio_generation.run_jobs()` 的 future wait 改为200ms短轮询，停止后及时取消未开始任务；提交/轮询请求返回后再次检查停止。下载链路新增可选 `should_stop`：requests 在流式数据块边界取消并关闭 response，curl 改为隐藏窗口 `Popen`，停止时有界 terminate/kill，重试等待可取消，停止后的本地文件不入资产库。 +- 2026-07-16：补齐真实 QThread 完成、终态信号丢失兜底、线程异常收尾、旧 token、重复停止、立即停止、排队 future 取消、下载取消 resume、requests response 关闭和 curl 子进程/临时文件清理测试。相关 90 项通过；隔离工作树全量 533 项 unittest、Ruff、compileall、`git diff --check` 全部通过。当前主工作区全量测试仅有3项原有封面默认模板 `papa1` 改名导致的失败，该未提交用户改动未纳入本任务。 diff --git a/tests/test_ai.py b/tests/test_ai.py index f339670..47bfa4a 100644 --- a/tests/test_ai.py +++ b/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: diff --git a/tests/test_image_studio_generation.py b/tests/test_image_studio_generation.py index 0daf0bb..91a7640 100644 --- a/tests/test_image_studio_generation.py +++ b/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) diff --git a/tests/test_product_suite_gui.py b/tests/test_product_suite_gui.py index c670b68..eecb8a4 100644 --- a/tests/test_product_suite_gui.py +++ b/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) diff --git a/tests/test_product_suite_workers.py b/tests/test_product_suite_workers.py index 7643a80..96283af 100644 --- a/tests/test_product_suite_workers.py +++ b/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)