diff --git a/app/gui/tabs/product_suite.py b/app/gui/tabs/product_suite.py index b418869..f763e6c 100644 --- a/app/gui/tabs/product_suite.py +++ b/app/gui/tabs/product_suite.py @@ -707,6 +707,15 @@ class SuiteTaskState: generation_terminal_streak: int = 0 pull_worker: object = None pull_thread: object = None + pull_run_token: str = "" + pull_stop_requested: bool = False + pull_cleanup_mode: str = "keep" + pull_before_asset_ids: set = field(default_factory=set) + pull_before_asset_states: list = field(default_factory=list) + pull_asset_ids: set = field(default_factory=set) + pull_download_asset_ids: set = field(default_factory=set) + pull_download_failed: int = 0 + pull_started_at: float = None import_worker: object = None import_thread: object = None import_created_draft: bool = False @@ -714,6 +723,7 @@ class SuiteTaskState: ai_thread: object = None download_queue: list = field(default_factory=list) downloads: dict = field(default_factory=dict) + download_tokens: dict = field(default_factory=dict) done: int = 0 failed: int = 0 total: int = 0 @@ -725,6 +735,9 @@ class SuiteTaskState: def generation_running(self): return self.worker is not None + def pull_running(self): + return bool(self.pull_run_token) + class ProductSuiteTab(QWidget): """Native PySide6 product-suite UI backed by image_studio services.""" @@ -757,6 +770,7 @@ class ProductSuiteTab(QWidget): self._displayed_state = None self._prompt_save_timers = {} self._generation_run_states = {} + self._pull_run_states = {} self._original_list_context = None self._loading = False self._result_refresh_pending = False @@ -1363,8 +1377,12 @@ class ProductSuiteTab(QWidget): state.worker.cancel() if state.ai_worker is not None: state.ai_worker.cancel() - if state.pull_worker is not None: - state.pull_worker.cancel() + if state.pull_running(): + state.pull_stop_requested = True + state.pull_cleanup_mode = "keep" + if state.pull_worker is not None: + state.pull_worker.cancel() + self._cancel_pull_downloads(state) if image_studio.is_draft_project(project): if draft_action == "delete": try: @@ -1385,6 +1403,8 @@ class ProductSuiteTab(QWidget): self._release_prompt_save_timer(state) if state.generation_run_token: self._generation_run_states.pop(state.generation_run_token, None) + if state.pull_run_token: + self._pull_run_states.pop(state.pull_run_token, None) self._retired_states.append(state) self._states.pop(state.key, None) self.task_tabs.removeTab(index) @@ -1602,6 +1622,13 @@ class ProductSuiteTab(QWidget): state.generation_job_ids = [] state.generation_mode = "batch" state.generation_retry_job_id = None + state.pull_run_token = "" + state.pull_stop_requested = False + state.pull_before_asset_ids = set() + state.pull_before_asset_states = [] + state.pull_asset_ids = set() + state.pull_download_asset_ids = set() + state.pull_download_failed = 0 state.done = state.failed = state.total = 0 state.started_at = None if state is self._displayed_state: @@ -1916,13 +1943,23 @@ class ProductSuiteTab(QWidget): def _update_context_actions(self, state): is_draft = self._is_draft_state(state) + pull_running = bool(state and state.pull_running()) self.pull_button.setEnabled( state is not None - and state.pull_worker is None - and (is_draft or self._valid_context(state, show_message=False)) + and ( + pull_running + or is_draft + or self._valid_context(state, show_message=False) + ) ) self.pull_button.setToolTip( - "需要先绑定正式商品ID" if is_draft else "拉取蝦皮主图" + "正在停止拉取" + if pull_running and state.pull_stop_requested + else "停止拉取蝦皮主图" + if pull_running + else "需要先绑定正式商品ID" + if is_draft + else "拉取蝦皮主图" ) self._refresh_item_id_hint(state) self._refresh_add_images_action(state) @@ -1982,6 +2019,10 @@ class ProductSuiteTab(QWidget): if run_token: thread.setProperty("productSuiteRunToken", run_token) thread.finished.connect(self._on_generation_thread_finished_signal) + pull_run_token = str(getattr(worker, "pull_run_token", "") or "") + if pull_run_token: + thread.setProperty("productSuitePullRunToken", pull_run_token) + thread.finished.connect(self._on_pull_thread_finished_signal) thread.start() return thread @@ -2158,6 +2199,9 @@ class ProductSuiteTab(QWidget): if state.generation_running(): self._status("生成中不能删除当前任务的商品原图", "warning") return + if state.pull_running(): + self._status("拉取蝦皮主图时不能删除当前任务的商品原图", "warning") + return normalized_ids = [] seen = set() for value in asset_ids or []: @@ -2243,6 +2287,9 @@ class ProductSuiteTab(QWidget): state = self._displayed_state if state is None: return + if state.pull_running(): + self._request_stop_pull(state) + return self._save_controls_to_state(state) if self._is_draft_state(state): self._message( @@ -2252,9 +2299,6 @@ class ProductSuiteTab(QWidget): return if not self._valid_context(state): return - if state.pull_worker is not None: - self._status("当前任务正在拉取蝦皮主图", "info") - return existing = [asset for asset in self._original_assets(state, include_missing=False) if _asset_usable(asset)] message = ( "店铺：%s\n" @@ -2276,49 +2320,406 @@ class ProductSuiteTab(QWidget): cancel_text="取消", ): return + if state.project_id is None: + project = image_studio.get_project_by_account_item( + state.account_alias, + state.item_id, + path=self.db_path, + ) + if project is not None: + state.project_id = int(project.id) + state.project_binding_state = project.binding_state + before_assets = ( + image_studio.list_assets( + state.project_id, + kind=image_studio.ASSET_KIND_ORIGINAL, + path=self.db_path, + ) + if state.project_id is not None + else [] + ) + pull_run_token = uuid.uuid4().hex worker = ImageStudioPullImagesWorker( state.account_alias, state.item_id, + pull_run_token=pull_run_token, db_path=self.db_path, config=self.config, ) state.pull_worker = worker - worker.finished.connect(lambda result, state=state: self._on_pull_finished(state, result)) - worker.cancelled.connect(lambda result, state=state: self._on_pull_finished(state, result)) - worker.failed.connect(lambda row, error, state=state: self._on_pull_failed(state, error)) + state.pull_run_token = pull_run_token + state.pull_stop_requested = False + state.pull_cleanup_mode = "keep" + state.pull_before_asset_ids = {int(asset.id) for asset in before_assets} + state.pull_before_asset_states = [ + { + "id": int(asset.id), + "status": asset.status, + "source_order": int(asset.source_order or 0), + } + for asset in before_assets + ] + state.pull_asset_ids = set() + state.pull_download_asset_ids = set() + state.pull_download_failed = 0 + state.pull_started_at = time.monotonic() + self._pull_run_states[pull_run_token] = state.key + self._log_pull_lifecycle( + state, + pull_run_token, + "started", + { + "before_count": len(state.pull_before_asset_ids), + }, + ) + worker.finished.connect( + lambda result, token=pull_run_token: self._on_pull_finished( + token, + result, + ) + ) + worker.cancelled.connect( + lambda result, token=pull_run_token: self._on_pull_finished( + token, + result, + ) + ) + worker.failed.connect( + lambda row, error, token=pull_run_token: self._on_pull_failed( + token, + error, + ) + ) state.pull_thread = self._start_thread(worker, "商品套图拉取蝦皮主图") if state is self._displayed_state: - self.pull_button.setText("正在拉取...") - self._update_context_actions(state) + self._apply_running_state(state) self._status("开始拉取蝦皮主图，可继续操作其他套图任务", "info") - def _on_pull_failed(self, state, error): + def _pull_state(self, pull_run_token): + token = str(pull_run_token or "") + state = self._states.get(self._pull_run_states.get(token)) + if state is None or state.pull_run_token != token: + return None + return state + + def _pull_stop_action(self): + box = QMessageBox(self) + box.setIcon(QMessageBox.Question) + box.setWindowTitle("停止拉取蝦皮主图") + box.setText("请选择停止后如何处理本轮已经拉取的图片。") + keep_button = box.addButton("停止并保留", QMessageBox.AcceptRole) + clear_button = box.addButton( + "停止并清除本次新增", + QMessageBox.DestructiveRole, + ) + continue_button = box.addButton("继续拉取", QMessageBox.RejectRole) + box.setDefaultButton(keep_button) + box.exec() + clicked = box.clickedButton() + if clicked is keep_button: + return "keep" + if clicked is clear_button: + return "clear_current" + if clicked is continue_button: + return "continue" + return "continue" + + def _request_stop_pull(self, state): + if state.pull_stop_requested: + self._status("正在停止当前拉取任务", "warning") + return + action = self._pull_stop_action() + if action == "continue": + return + state.pull_stop_requested = True + state.pull_cleanup_mode = action + if state.pull_worker is not None: + state.pull_worker.cancel() + self._cancel_pull_downloads(state) + self._log_pull_lifecycle( + state, + state.pull_run_token, + "stop_requested", + {"cleanup_mode": action}, + level="WARNING", + ) + if state is self._displayed_state: + self._apply_running_state(state) + self._status("已请求停止拉取蝦皮主图", "warning") + self._maybe_finalize_pull(state, state.pull_run_token) + + def _cancel_pull_downloads(self, state): + pull_ids = set(state.pull_download_asset_ids) + if not pull_ids: + return + state.download_queue = [ + asset_id + for asset_id in state.download_queue + if int(asset_id) not in pull_ids + ] + for asset_id in list(pull_ids): + running = state.downloads.get(int(asset_id)) + if running is None: + state.pull_download_asset_ids.discard(int(asset_id)) + state.download_tokens.pop(int(asset_id), None) + continue + worker, _ = running + worker.cancel() + + def _on_pull_failed(self, pull_run_token, error): + state = self._pull_state(pull_run_token) + if state is None or state.pull_stop_requested: + return self._status("拉取蝦皮主图失败：%s" % _user_error(error), "danger") - def _on_pull_finished(self, state, result): + def _on_pull_finished(self, pull_run_token, result): + state = self._pull_state(pull_run_token) + if state is None: + return state.pull_worker = None state.pull_thread = None - if state is self._displayed_state: - self.pull_button.setText("拉取蝦皮主图") - if result.get("ok") is False: - self._message("拉取蝦皮主图失败", _user_error(result.get("error"))) + result = dict(result or {}) + if result.get("ok") is False and not state.pull_stop_requested: + self._finalize_pull( + state, + pull_run_token, + error=_user_error(result.get("error")), + ) + return + project = result.get("project") + if project is not None: + state.project_id = int(project.id) + state.project_binding_state = project.binding_state + assets = [ + asset + for asset in (result.get("assets") or []) + if asset.status != image_studio.ASSET_STATUS_MISSING + ][: image_studio_images.MAX_ORIGINAL_ASSETS] + state.pull_asset_ids.update( + int(asset.id) + for asset in assets + if getattr(asset, "remote_url", None) + ) + if result.get("cancelled") or state.pull_stop_requested: + self._cancel_pull_downloads(state) else: - project = result.get("project") - if project is not None: - state.project_id = int(project.id) - state.project_binding_state = project.binding_state - assets = [ - asset - for asset in (result.get("assets") or []) - if asset.status != image_studio.ASSET_STATUS_MISSING - ][: image_studio_images.MAX_ORIGINAL_ASSETS] - self._queue_original_downloads(state, assets) + self._queue_original_downloads( + state, + assets, + pull_run_token=pull_run_token, + ) self._status("已读取%d张蝦皮主图，正在后台下载" % len(assets), "success") if state is self._displayed_state: self._refresh_originals(state) - self._update_context_actions(state) + self._apply_running_state(state) + self._maybe_finalize_pull(state, pull_run_token) - def _queue_original_downloads(self, state, assets): + @Slot() + def _on_pull_thread_finished_signal(self): + sender = self.sender() + token = str( + sender.property("productSuitePullRunToken") + if sender is not None + else "" + ) + QTimer.singleShot(0, lambda token=token: self._handle_pull_thread_finished(token)) + + def _handle_pull_thread_finished(self, pull_run_token): + state = self._pull_state(pull_run_token) + if state is None or state.pull_worker is None: + return + state.pull_worker = None + state.pull_thread = None + if state.pull_stop_requested: + self._maybe_finalize_pull(state, pull_run_token) + else: + self._finalize_pull( + state, + pull_run_token, + error="拉取线程已结束，请稍后重试", + ) + + def _maybe_finalize_pull(self, state, pull_run_token): + if self._pull_state(pull_run_token) is not state: + return False + if state.pull_worker is not None or state.pull_download_asset_ids: + return False + return self._finalize_pull(state, pull_run_token) + + def _clear_current_pull_assets(self, state): + new_ids = sorted( + set(state.pull_asset_ids) - set(state.pull_before_asset_ids) + ) + removable = [] + retained = 0 + for asset_id in new_ids: + asset = image_studio.get_asset(asset_id, path=self.db_path) + if ( + asset is None + or int(asset.project_id) != int(state.project_id or 0) + or asset.kind != image_studio.ASSET_KIND_ORIGINAL + or not asset.remote_url + ): + continue + if image_studio.asset_reference_counts( + asset.id, + path=self.db_path, + )["total"]: + retained += 1 + else: + removable.append(asset.id) + removed = 0 + if removable: + image_studio.remove_original_assets_if_unused( + state.project_id, + removable, + path=self.db_path, + ) + removed = len(removable) + existing_ids = { + int(asset.id) + for asset in image_studio.list_assets( + state.project_id, + kind=image_studio.ASSET_KIND_ORIGINAL, + path=self.db_path, + ) + } + snapshot = [ + item + for item in state.pull_before_asset_states + if int(item["id"]) in existing_ids + ] + if snapshot: + image_studio.restore_original_asset_snapshot( + state.project_id, + snapshot, + path=self.db_path, + ) + return removed, retained + + def _finalize_pull(self, state, pull_run_token, *, error=""): + if self._pull_state(pull_run_token) is not state: + return False + cancelled = state.pull_stop_requested + cleanup_mode = state.pull_cleanup_mode + pulled_count = len(state.pull_asset_ids) + new_count = len(set(state.pull_asset_ids) - set(state.pull_before_asset_ids)) + removed = 0 + retained = 0 + cleanup_error = "" + if cancelled and cleanup_mode == "clear_current" and state.project_id is not None: + try: + removed, retained = self._clear_current_pull_assets(state) + except Exception as exc: + cleanup_error = _user_error(exc) + elapsed = int( + max(0, time.monotonic() - state.pull_started_at) + if state.pull_started_at + else 0 + ) + self._pull_run_states.pop(pull_run_token, None) + state.pull_worker = None + state.pull_thread = None + state.pull_run_token = "" + state.pull_stop_requested = False + state.pull_cleanup_mode = "keep" + state.pull_before_asset_ids = set() + state.pull_before_asset_states = [] + state.pull_asset_ids = set() + state.pull_download_asset_ids = set() + failed = int(state.pull_download_failed or 0) + self._log_pull_lifecycle( + state, + pull_run_token, + "finalized", + { + "cancelled": cancelled, + "cleanup_mode": cleanup_mode, + "pulled_count": pulled_count, + "new_count": new_count, + "removed_count": removed, + "retained_count": retained, + "failed_count": failed, + "elapsed_seconds": elapsed, + "has_error": bool(error or cleanup_error), + }, + level="WARNING" if cancelled or error or cleanup_error else "INFO", + ) + state.pull_download_failed = 0 + state.pull_started_at = None + if state is self._displayed_state: + self._refresh_originals(state) + self._apply_running_state(state) + if cleanup_error: + self._message( + "停止拉取后清理失败", + "拉取已停止，但本轮图片清理失败：%s" % cleanup_error, + ) + self._status("拉取已停止，本轮图片清理失败", "danger") + return True + if error: + if state is self._displayed_state: + self._message("拉取蝦皮主图失败", error) + self._status("拉取蝦皮主图失败：%s" % error, "danger") + return True + if cancelled: + kept = max(0, new_count - removed) + if state is self._displayed_state: + self._message( + "拉取蝦皮主图已停止", + "本轮新增%d张：保留%d张，清理%d张，" + "因引用保留%d张，下载失败%d张；总用时%d秒。" + % ( + new_count, + kept, + removed, + retained, + failed, + elapsed, + ), + icon=QMessageBox.Information, + ) + self._status( + "拉取已停止：保留%d张，清理%d张" % (kept, removed), + "warning", + ) + return True + self._status( + "蝦皮主图拉取完成：本轮%d张，下载失败%d张，用时%d秒" + % (pulled_count, failed, elapsed), + "success" if not failed else "warning", + ) + return True + + def _log_pull_lifecycle( + self, + state, + pull_run_token, + event, + payload=None, + *, + level="INFO", + ): + data = { + "pull_run_token": str(pull_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_pull", + 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 _queue_original_downloads(self, state, assets, *, pull_run_token=None): if state is None: return queued = set(state.download_queue) | set(state.downloads) @@ -2328,6 +2729,9 @@ class ProductSuiteTab(QWidget): asset_id = int(asset.id) if asset_id not in queued: state.download_queue.append(asset_id) + state.download_tokens[asset_id] = str(pull_run_token or "") + if pull_run_token: + state.pull_download_asset_ids.add(asset_id) queued.add(asset_id) self._start_queued_downloads(state) @@ -2340,14 +2744,15 @@ class ProductSuiteTab(QWidget): config=self.config, max_retries=2, ) + download_token = str(state.download_tokens.get(asset_id, "") or "") worker.finished.connect( - lambda result, state=state, asset_id=asset_id: self._on_download_finished( - state, asset_id, result + lambda result, state=state, asset_id=asset_id, token=download_token: self._on_download_finished( + state, asset_id, result, pull_run_token=token ) ) worker.cancelled.connect( - lambda result, state=state, asset_id=asset_id: self._on_download_finished( - state, asset_id, result + lambda result, state=state, asset_id=asset_id, token=download_token: self._on_download_finished( + state, asset_id, result, pull_run_token=token ) ) worker.failed.connect( @@ -2359,9 +2764,24 @@ class ProductSuiteTab(QWidget): thread = self._start_thread(worker, "商品套图下载原图") state.downloads[asset_id] = (worker, thread) - def _on_download_finished(self, state, asset_id, result): + def _on_download_finished( + self, + state, + asset_id, + result, + *, + pull_run_token="", + ): + expected_token = str(state.download_tokens.get(int(asset_id), "") or "") + if expected_token != str(pull_run_token or ""): + return state.downloads.pop(int(asset_id), None) + state.download_tokens.pop(int(asset_id), None) + if pull_run_token: + state.pull_download_asset_ids.discard(int(asset_id)) if result.get("ok") is False: + if pull_run_token: + state.pull_download_failed += 1 self._status( "商品原图 #%d 下载失败：%s" % (asset_id, _user_error(result.get("error"))), "danger", @@ -2371,6 +2791,8 @@ class ProductSuiteTab(QWidget): if state is self._displayed_state: self._refresh_originals(state) self._start_queued_downloads(state) + if pull_run_token: + self._maybe_finalize_pull(state, pull_run_token) def _rebuild_categories(self, state): while self.category_grid.count(): @@ -3221,9 +3643,16 @@ class ProductSuiteTab(QWidget): def _apply_running_state(self, state): generation_running = state.generation_running() - self.pull_button.setText("正在拉取..." if state.pull_worker is not None else "拉取蝦皮主图") - self.account_combo.setEnabled(not generation_running and state.pull_worker is None) - self.item_id_edit.setEnabled(not generation_running and state.pull_worker is None) + pull_running = state.pull_running() + self.pull_button.setText( + "正在停止..." + if pull_running and state.pull_stop_requested + else "停止拉取蝦皮" + if pull_running + else "拉取蝦皮主图" + ) + self.account_combo.setEnabled(not generation_running and not pull_running) + self.item_id_edit.setEnabled(not generation_running and not pull_running) self._refresh_add_images_action(state) self.original_list.setEnabled(not generation_running) self._refresh_original_selection_controls() @@ -3522,4 +3951,5 @@ class ProductSuiteTab(QWidget): for state in list(self._states.values()) + list(self._retired_states): self._release_prompt_save_timer(state) self._generation_run_states.clear() + self._pull_run_states.clear() super().closeEvent(event) diff --git a/app/gui/workers.py b/app/gui/workers.py index 2ce1d5d..38d006f 100644 --- a/app/gui/workers.py +++ b/app/gui/workers.py @@ -90,31 +90,58 @@ def _image_studio_user_detail(detail): class ImageStudioPullImagesWorker(BaseWorker): """Read Shopee main image URLs for one AI studio project in background.""" - def __init__(self, account_alias, item_id, *, db_path=None, config=None): + def __init__( + self, + account_alias, + item_id, + *, + pull_run_token="", + db_path=None, + config=None, + ): super().__init__() self.account_alias = account_alias self.item_id = item_id + self.pull_run_token = str(pull_run_token or "") self.db_path = db_path self.config = config def execute(self): + if self.should_cancel(): + return { + "pull_run_token": self.pull_run_token, + "cancelled": True, + "assets": [], + } + def on_step(payload): self.log.emit(_format_image_studio_event(payload)) - result = image_studio.pull_remote_main_image_urls( - self.account_alias, - self.item_id, - path=self.db_path, - config=self.config, - on_step=on_step, - ) + try: + result = image_studio.pull_remote_main_image_urls( + self.account_alias, + self.item_id, + path=self.db_path, + config=self.config, + on_step=on_step, + should_stop=self.should_cancel, + ) + except image_studio.ImageStudioPullCancelled as exc: + return { + "pull_run_token": self.pull_run_token, + "cancelled": True, + "project": exc.project, + "assets": list(exc.assets or []), + } project = result.get("project") assets = result.get("assets") or [] return { + "pull_run_token": self.pull_run_token, "project": project, "assets": assets, "count": len(assets), "account": result.get("account"), + "cancelled": self.should_cancel(), } @@ -158,7 +185,10 @@ class ImageStudioDownloadOriginalWorker(BaseWorker): self.asset_id, path=self.db_path, config=self.config, + should_stop=self.should_cancel, ) + except image_studio_images.ImageStudioImageCancelled: + return {"asset_id": self.asset_id, "cancelled": True} except Exception: retry = attempt if attempt >= attempts: diff --git a/app/image_studio.py b/app/image_studio.py index b847fd2..441be40 100644 --- a/app/image_studio.py +++ b/app/image_studio.py @@ -116,6 +116,15 @@ class ImageStudioProjectConflictError(ImageStudioError): """Raised when a draft cannot be bound because the formal project already exists.""" +class ImageStudioPullCancelled(ImageStudioError): + """Raised when a read-only Shopee image pull stops at a safe boundary.""" + + def __init__(self, project=None, assets=None): + super().__init__("用户停止拉取蝦皮主图") + self.project = project + self.assets = list(assets or []) + + def _now() -> str: return datetime.now().isoformat(timespec="seconds") @@ -825,6 +834,93 @@ def remove_original_assets_if_unused(project_id, asset_ids, path=None, conn=None return [_row_to_dataclass(by_id[asset_id], ImageStudioAsset) for asset_id in ordered_ids] +def restore_original_asset_snapshot(project_id, snapshots, path=None, conn=None): + """Restore status/order for original assets that existed before a pull.""" + + project_id = int(project_id) + normalized = [] + seen = set() + for value in snapshots or []: + asset_id = int(_get(value, "id")) + if asset_id in seen: + continue + status = str(_get(value, "status") or ASSET_STATUS_AVAILABLE) + if status not in ASSET_STATUSES: + raise db.DbError("商品原图快照状态无效") + normalized.append( + { + "id": asset_id, + "status": status, + "source_order": max(0, int(_get(value, "source_order") or 0)), + } + ) + seen.add(asset_id) + if not normalized: + return [] + ids = [item["id"] for item in normalized] + placeholders = ",".join("?" for _ in ids) + with _connection(conn, path) as database: + with database: + rows = database.execute( + f""" + SELECT id FROM image_studio_assets + WHERE project_id = ? AND kind = ? AND id IN ({placeholders}) + """, + [project_id, ASSET_KIND_ORIGINAL, *ids], + ).fetchall() + if {int(row["id"]) for row in rows} != set(ids): + raise db.DbError("拉取前商品原图快照已失效") + now = _now() + for item in normalized: + database.execute( + """ + UPDATE image_studio_assets + SET status = ?, source_order = ?, updated_at = ? + WHERE id = ? AND project_id = ? AND kind = ? + """, + ( + item["status"], + item["source_order"], + now, + item["id"], + project_id, + ASSET_KIND_ORIGINAL, + ), + ) + remaining_rows = database.execute( + f""" + SELECT id FROM image_studio_assets + WHERE project_id = ? AND kind = ? + AND id NOT IN ({placeholders}) + ORDER BY source_order, id + """, + [project_id, ASSET_KIND_ORIGINAL, *ids], + ).fetchall() + next_order = max( + [int(item["source_order"]) for item in normalized] + [0] + ) + 1 + for row in remaining_rows: + database.execute( + """ + UPDATE image_studio_assets + SET source_order = ?, updated_at = ? + WHERE id = ? AND project_id = ? AND kind = ? + """, + ( + next_order, + now, + int(row["id"]), + project_id, + ASSET_KIND_ORIGINAL, + ), + ) + next_order += 1 + return [ + get_asset(item["id"], conn=database) + for item in normalized + ] + + def sync_original_asset_urls(project_id, image_urls, path=None, conn=None, max_assets=16): """Store the read-only Shopee main image URL snapshot as remote-only assets.""" @@ -1183,10 +1279,14 @@ def pull_remote_main_image_urls( config=None, login_timeout=8, on_step=None, + should_stop=None, ): """Create/open an AI studio project and read Shopee main image URLs without editing.""" cfg = appconfig.load_config() if config is None else config + should_stop = should_stop or (lambda: False) + if should_stop(): + raise ImageStudioPullCancelled() database_path = _db_path(path, cfg) db.init_db(database_path) item = _normalize_item_id(item_id) @@ -1196,6 +1296,8 @@ def pull_remote_main_image_urls( raise ImageStudioError(f"AI工场账号不可用：{exc}") from exc project = create_or_get_project(account, item_id=item, path=database_path) + if should_stop(): + raise ImageStudioPullCancelled(project=project) readiness = _ensure_account_ready_for_read( account, path=database_path, @@ -1203,16 +1305,26 @@ def pull_remote_main_image_urls( login_timeout=login_timeout, on_step=on_step, ) + if should_stop(): + raise ImageStudioPullCancelled(project=project) cdp = None try: + if should_stop(): + raise ImageStudioPullCancelled(project=project) _notify_step(on_step, "open_product", "start", f"商品 {item}") cdp = editor.open_product(account, item, on_step=on_step, bring_to_front=False) + if should_stop(): + raise ImageStudioPullCancelled(project=project) _notify_step(on_step, "open_product", "success", f"商品 {item}") _notify_step(on_step, "read_main_images", "start", f"商品 {item}") images = editor.read_product_image_urls(cdp) + if should_stop(): + raise ImageStudioPullCancelled(project=project) if not images: raise ImageStudioError("未读取到蝦皮商品主图 URL") assets = sync_original_asset_urls(project.id, images, path=database_path) + if should_stop(): + raise ImageStudioPullCancelled(project=project, assets=assets) _notify_step(on_step, "read_main_images", "success", f"读取 {len(images)} 张主图 URL") project = get_project(project.id, path=database_path) return { diff --git a/app/image_studio_images.py b/app/image_studio_images.py index bf61fe0..071cf43 100644 --- a/app/image_studio_images.py +++ b/app/image_studio_images.py @@ -32,6 +32,10 @@ class ImageStudioImageError(RuntimeError): """Raised when a remote image cannot be safely loaded or saved.""" +class ImageStudioImageCancelled(ImageStudioImageError): + """Raised when an original image download is cancelled safely.""" + + @dataclass(frozen=True) class RemoteImage: url: str @@ -127,10 +131,14 @@ def download_remote_image( max_bytes=ORIGINAL_MAX_BYTES, timeout=(DEFAULT_CONNECT_TIMEOUT_SECONDS, DEFAULT_READ_TIMEOUT_SECONDS), session=None, + should_stop=None, ) -> RemoteImage: """Download remote image bytes with SSRF, timeout and size guards.""" url = str(url or "").strip() + should_stop = should_stop or (lambda: False) + if should_stop(): + raise ImageStudioImageCancelled("用户停止下载商品原图") _assert_public_http_url(url) client = session or _session() try: @@ -142,38 +150,55 @@ def download_remote_image( ) except requests.exceptions.RequestException as exc: raise ImageStudioImageError(f"远程图片下载失败：{exc}") from exc - status = getattr(response, "status_code", 200) - if status >= 400: - raise ImageStudioImageError(f"远程图片下载失败：HTTP {status}") - final_url = str(getattr(response, "url", "") or url) - _assert_public_http_url(final_url) - content_type = _validate_content_type(getattr(response, "headers", {}).get("Content-Type", "")) - content_length = getattr(response, "headers", {}).get("Content-Length") - if content_length: - try: - if int(content_length) > int(max_bytes): + try: + if should_stop(): + raise ImageStudioImageCancelled("用户停止下载商品原图") + status = getattr(response, "status_code", 200) + if status >= 400: + raise ImageStudioImageError(f"远程图片下载失败：HTTP {status}") + final_url = str(getattr(response, "url", "") or url) + _assert_public_http_url(final_url) + content_type = _validate_content_type( + getattr(response, "headers", {}).get("Content-Type", "") + ) + content_length = getattr(response, "headers", {}).get("Content-Length") + if content_length: + try: + if int(content_length) > int(max_bytes): + raise ImageStudioImageError("远程图片超过大小上限") + except ValueError: + pass + chunks = [] + total = 0 + iterator = ( + response.iter_content(chunk_size=65536) + if hasattr(response, "iter_content") + else [response.content] + ) + for chunk in iterator: + if should_stop(): + raise ImageStudioImageCancelled("用户停止下载商品原图") + if not chunk: + continue + total += len(chunk) + if total > int(max_bytes): raise ImageStudioImageError("远程图片超过大小上限") - except ValueError: - pass - 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 > int(max_bytes): - raise ImageStudioImageError("远程图片超过大小上限") - chunks.append(chunk) - content = b"".join(chunks) - _image_info(content) - return RemoteImage( - url=url, - content=content, - content_type=content_type, - final_url=final_url, - redirected=final_url != url, - ) + chunks.append(chunk) + if should_stop(): + raise ImageStudioImageCancelled("用户停止下载商品原图") + content = b"".join(chunks) + _image_info(content) + return RemoteImage( + url=url, + content=content, + content_type=content_type, + final_url=final_url, + redirected=final_url != url, + ) + finally: + close = getattr(response, "close", None) + if callable(close): + close() def _image_info(image_bytes): @@ -509,10 +534,21 @@ def restore_trashed_asset(record, *, path=None, config=None): raise ImageStudioImageError(f"撤销删除生成图片失败：{exc}") from exc -def download_original_asset(asset_id, *, path=None, config=None, image_root=None, session=None): +def download_original_asset( + asset_id, + *, + path=None, + config=None, + image_root=None, + session=None, + should_stop=None, +): """Download one Shopee original image into originals/ and mark its asset available.""" cfg = appconfig.load_config() if config is None else config + should_stop = should_stop or (lambda: False) + if should_stop(): + raise ImageStudioImageCancelled("用户停止下载商品原图") database_path = path or appconfig.db_path(cfg) asset = image_studio.get_asset(asset_id, path=database_path) if asset is None: @@ -525,15 +561,26 @@ def download_original_asset(asset_id, *, path=None, config=None, image_root=None project = image_studio.get_project(asset.project_id, path=database_path) if project is None: raise ImageStudioImageError("原图资产所属项目不存在") - remote = download_remote_image(asset.remote_url, max_bytes=ORIGINAL_MAX_BYTES, session=session) + remote = download_remote_image( + asset.remote_url, + max_bytes=ORIGINAL_MAX_BYTES, + session=session, + should_stop=should_stop, + ) + if should_stop(): + raise ImageStudioImageCancelled("用户停止下载商品原图") info = _image_info(remote.content) directory, stem = _original_file_path(project, asset, image_root or appconfig.image_dir(cfg)) os.makedirs(directory, exist_ok=True) final_path = os.path.join(directory, stem + _extension_for_format(info["format"])) temp_path = final_path + ".tmp-" + uuid.uuid4().hex try: + if should_stop(): + raise ImageStudioImageCancelled("用户停止下载商品原图") with open(temp_path, "wb") as fh: fh.write(remote.content) + if should_stop(): + raise ImageStudioImageCancelled("用户停止下载商品原图") with open(temp_path, "rb") as fh: _image_info(fh.read()) os.replace(temp_path, final_path) diff --git a/docs/04-architecture.md b/docs/04-architecture.md index a383aa6..592077f 100644 --- a/docs/04-architecture.md +++ b/docs/04-architecture.md @@ -438,6 +438,7 @@ data/images///__new. # AI 生成的新 - `image_studio_projects.suite_settings_json` 持久化套图设置，旧数据库由 `db.init_db()` 原位补列，默认 `{}`；`draft_prompt` 继续保存卖点文本。`image_studio_assets` 中有效商品原图最多16张，历史 missing 记录不占有效名额；手工原图不会因再次同步蝦皮 URL 被误标 missing。⑥原图列表的批量勾选只保存在当前 `SuiteTaskState` 对应的界面上下文，不写库；批量移除由 `remove_original_assets_if_unused()` 一次校验项目归属、原图类型和 job/终选引用，并在单个 SQLite 事务中删除资产行、连续重排 `source_order`。服务不删除本地文件或蝦皮线上图片，任一资产校验失败时整批回滚。 - T-636 起，`image_studio_projects` 增加 `binding_state`（`draft` / `bound`）和稳定 `storage_key`。既有项目迁移为 `bound`，并以原 `item_id` 回填 `storage_key`；项目目录改用 `storage_key`，因此临时草稿绑定正式商品 ID 后不移动目录、不改写已有资产路径。草稿内部使用 `draft_` 作为仅数据库可见的非空 `item_id`，GUI 输入框始终保持空白，用户日志和 cmhub 提示词只使用“临时草稿”或“未绑定商品”，不得暴露该内部值。 - ⑥已选账号但未填写商品 ID 时允许输入卖点、导入、拖入或粘贴本地图片。非空卖点在现有防抖稳定后会创建一个可恢复临时草稿并保存到 `image_studio_projects.draft_prompt`；首次有效图片导入也会创建草稿。空白卖点、取消选择和全部导入失败不保留空草稿。只有卖点的草稿同样属于可恢复业务内容；卖点清空且没有资产/job 时仍可按空草稿规则清理。草稿可管理本地图片、AI 帮写、生成套图、查看历史和打开结果目录，但在创建 worker、启动 Chrome 或执行 CDP 前禁止「拉取蝦皮主图」。输入合法数字商品 ID 后，经确认原地绑定同一个 `project_id`；资产、job、selection、提示词、套图设置和 `storage_key` 均保持不变。若同账号目标 ID（含软删除项目）已存在则拒绝覆盖或合并。 +- ⑥拉取蝦皮主图使用独立内存 `pull_run_token` 覆盖 URL 读取和本轮后台下载。运行中按钮提供「停止并保留 / 停止并清除本次新增 / 继续拉取」；停止采用协作式安全边界，不强杀线程、Chrome 或已发出的 CDP/HTTP 请求。选择清除时只移除本轮新增、属于当前项目且未被 job/selection 引用的远程原图，并恢复拉取前已有原图的状态和顺序；本地手动导入图片、拉取前已有图片、用户文件和蝦皮线上图片不删除。旧 token 的 URL/下载迟到结果不得覆盖新一轮状态。 - 启动时恢复未软删除、至少含一条资产或生成任务的草稿为独立中文“临时草稿”标签，按最近更新时间排序。关闭非空草稿可选择保留、软删除或取消；软删除不物理删除图片目录。③「更新蝦皮」只处理正式任务，不接受临时草稿。 - 第六 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/api.md b/docs/api.md index 15cde2c..681af8a 100644 --- a/docs/api.md +++ b/docs/api.md @@ -378,12 +378,14 @@ create_job(project_id, source_asset_id=None, job_type="main", prompt="", ...) -> list_jobs(project_id, statuses=None, path=None) -> list[ImageStudioJob] list_resumable_jobs(project_id=None, include_failed_downloads=False, path=None) -> list[ImageStudioJob] replace_selections(project_id, selection_type, asset_ids, path=None) -> list[ImageStudioSelection] -pull_remote_main_image_urls(account_or_alias, item_id, path=None, config=None) -> dict +pull_remote_main_image_urls(account_or_alias, item_id, path=None, config=None, + should_stop=None) -> dict # app/image_studio_images.py download_remote_image(url, max_bytes=..., timeout=(connect, read)) -> RemoteImage load_thumbnail(url, key=None, max_size=220) -> ThumbnailResult -download_original_asset(asset_id, path=None, config=None) -> ImageStudioAsset +download_original_asset(asset_id, path=None, config=None, + should_stop=None) -> ImageStudioAsset import_original_files(project_id, file_paths, path=None, config=None) -> dict import_original_bytes(project_id, content, filename_hint="clipboard.png", ...) -> ImageStudioAsset trash_generated_asset(asset_id, path=None, config=None) -> dict @@ -433,7 +435,7 @@ class GenerateWorker(BaseWorker) # ② 后台生成：ai.generate_batc class ApplyWorker(BaseWorker) # ③ 后台更新：账号就绪预检 -> 检查或按批调用 editor.apply_task(...) -> db.set_applied/mark_skipped class WriteBackWorker(BaseWorker) # ①/③ 后台回写：旧字段或更新结果写回原 Excel class AIModelTestWorker(BaseWorker) # ⑤ 后台测试 AI 模型连接：appconfig.test_ai_model -class ImageStudioPullImagesWorker(BaseWorker) # ⑥ 后台只读拉蝦皮原主图 URL +class ImageStudioPullImagesWorker(BaseWorker) # ⑥ 后台只读拉蝦皮原主图 URL；安全边界协作停止 class ImageStudioDownloadOriginalWorker(BaseWorker)# ⑥ 后台下载远程原图 class ImageStudioGenerateJobsWorker(BaseWorker) # ⑥ 后台提交/查询/下载 cmhub 生图 job class ImageStudioResumeJobsWorker(BaseWorker) # ⑥ 后台恢复已有 task_id 的生图 job diff --git a/docs/routes.md b/docs/routes.md index 89880b2..f5979ec 100644 --- a/docs/routes.md +++ b/docs/routes.md @@ -243,6 +243,6 @@ | `ApplyWorker(BaseWorker)` | ③ | 账号就绪预检、检查本轮更新、按每批最大条数分批、按账号并行或串行调用 `editor.apply_task(...)`、逐条 `set_applied()`，失败继续，写运行日志 | | `AIModelTestWorker(BaseWorker)` | ⑤ | 后台调用 `appconfig.test_ai_model()` 测试模型连接 | | `WriteBackWorker(BaseWorker)` | ①③ | ①回写旧字段；③回写新标题/新封面/更新状态 | -| `ImageStudioPullImagesWorker / ImageStudioDownloadOriginalWorker / ProductSuiteImportImagesWorker / ProductSuiteGenerateWorker / ProductSuiteAiWriteWorker` | ⑥ | 后台执行只读拉主图、远程原图下载、本地图片校验复制、cmhub 套图生成与AI帮写；不直接操作 QWidget | +| `ImageStudioPullImagesWorker / ImageStudioDownloadOriginalWorker / ProductSuiteImportImagesWorker / ProductSuiteGenerateWorker / ProductSuiteAiWriteWorker` | ⑥ | 后台执行只读拉主图、远程原图下载、本地图片校验复制、cmhub 套图生成与AI帮写；拉图和本轮下载支持安全边界协作停止，worker 不直接操作 QWidget | > 采集、生成、更新都是耗时操作，使用 `QObject` worker + `QThread`。Worker 不直接操作 QWidget，只通过 signal 通知主线程刷新 UI。 diff --git a/docs/tasks/T-642.md b/docs/tasks/T-642.md index 6aa8da4..9f9bb69 100644 --- a/docs/tasks/T-642.md +++ b/docs/tasks/T-642.md @@ -3,7 +3,7 @@ id: T-642 title: 商品套图拉取主图可停止与本轮图片清理 phase: 7 deps: [T-641] -status: TODO +status: DONE created: 2026-07-16 --- @@ -151,4 +151,8 @@ git diff --check ## 执行记录 -- 待执行。 +- 2026-07-16：⑥拉取主图新增内存 `pull_run_token` 生命周期，覆盖 URL 读取和本轮远程原图下载。按钮在活动期间显示「停止拉取蝦皮」，停止请求后显示「正在停止...」；停止弹窗提供「停止并保留 / 停止并清除本次新增 / 继续拉取」。重复停止不重复弹窗，账号和商品 ID 在本轮结束前保持锁定，其他任务仍可切换使用。 +- 2026-07-16：`ImageStudioPullImagesWorker` 和 `pull_remote_main_image_urls()` 增加可选 `should_stop`，在账号准备、打开商品页、读取 URL 和同步资产前后检查；当前 CDP 请求返回后安全停止，并继续复用 `close_readonly_product()` 关闭本轮只读 tab。线程结束增加 token 校验和幂等兜底；旧轮次 URL/下载迟到结果不会清除新 worker 或覆盖新状态。 +- 2026-07-16：本轮下载使用 asset ID → token 归属，只取消本轮排队和活动下载，不影响用户此前手动触发的下载。远程图片 requests 下载在请求返回、分块读取和保存前检查取消并关闭 response、清理临时文件；不使用线程强杀或结束 Chrome。 +- 2026-07-16：「停止并清除本次新增」按拉取前 asset 快照与本轮远程 asset 差集处理，只移除当前项目中未被 job/selection 引用的新远程原图；被引用图片保留。本地手动导入和拉取前图片不删除，旧图片状态/顺序恢复，新但被引用的图片追加到旧图片之后；本地文件和蝦皮线上图片不删除。 +- 2026-07-16：拉取生命周期诊断日志只记录 token 短值、项目/商品、数量、清理策略和用时。同步更新架构、API 与路由文档；补齐安全边界取消、response 关闭、快照恢复、本轮下载隔离、按钮三态、清理/引用保护、旧 token、线程结束兜底和真实 `QThread` 完成回归测试。相关 81 项通过；商品套图 GUI 42 项在主工作区通过，其余测试文件在短路径隔离 worktree 中逐文件通过，共覆盖 550 项；Ruff、`compileall` 与 `git diff --check` 均通过。未连接真实账号执行 CDP 停止冒烟，需在测试商品 `51100639510` 上人工确认当前有界请求返回后的停止时延和自动新建 tab 关闭表现。 diff --git a/tests/test_image_studio.py b/tests/test_image_studio.py index 0383eef..17369ad 100644 --- a/tests/test_image_studio.py +++ b/tests/test_image_studio.py @@ -1060,6 +1060,95 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase): self.assert_removed(temp_dir) + def test_pull_remote_main_image_urls_stops_after_open_and_closes_created_tab(self): + with self.make_temp_dir() as temp_dir: + cfg = self._config(temp_dir) + accounts.create_account("主店", "alias-a", debug_port=9222, config=cfg) + cdp = SimpleNamespace() + stop_values = iter([False, False, False, False, True]) + + with mock.patch("app.image_studio.chrome.is_running", return_value=True), \ + mock.patch( + "app.image_studio.accounts.detect_login", + return_value={"logged_in": True}, + ), \ + mock.patch( + "app.image_studio.editor.open_product", + return_value=cdp, + ) as open_product, \ + mock.patch( + "app.image_studio.editor.read_product_image_urls", + ) as read_urls, \ + mock.patch( + "app.image_studio.editor.close_readonly_product", + ) as close_readonly: + with self.assertRaises(image_studio.ImageStudioPullCancelled): + image_studio.pull_remote_main_image_urls( + "alias-a", + "51100639510", + path=cfg["db_path"], + config=cfg, + should_stop=lambda: next(stop_values), + ) + + open_product.assert_called_once() + read_urls.assert_not_called() + close_readonly.assert_called_once_with(cdp) + + self.assert_removed(temp_dir) + + def test_restore_original_asset_snapshot_restores_existing_state(self): + with self.make_temp_dir() as temp_dir: + db_path = os.path.join(temp_dir, "cmshopee.db") + db.init_db(db_path) + project = image_studio.create_or_get_project( + account_alias="alias-a", + account_slug="alias_a", + item_id="51100639510", + path=db_path, + ) + assets = image_studio.sync_original_asset_urls( + project.id, + [ + {"index": 1, "src": "https://susercontent.com/one.jpg"}, + {"index": 2, "src": "https://susercontent.com/two.jpg"}, + ], + path=db_path, + ) + snapshot = [ + { + "id": assets[0].id, + "status": image_studio.ASSET_STATUS_AVAILABLE, + "source_order": 1, + }, + { + "id": assets[1].id, + "status": image_studio.ASSET_STATUS_AVAILABLE, + "source_order": 2, + }, + ] + image_studio.sync_original_asset_urls( + project.id, + [{"index": 1, "src": "https://susercontent.com/two.jpg"}], + path=db_path, + ) + + restored = image_studio.restore_original_asset_snapshot( + project.id, + snapshot, + path=db_path, + ) + + self.assertEqual( + [ + (image_studio.ASSET_STATUS_AVAILABLE, 1), + (image_studio.ASSET_STATUS_AVAILABLE, 2), + ], + [(asset.status, asset.source_order) for asset in restored], + ) + + self.assert_removed(temp_dir) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_image_studio_images.py b/tests/test_image_studio_images.py index 31dff79..9ff0b91 100644 --- a/tests/test_image_studio_images.py +++ b/tests/test_image_studio_images.py @@ -29,6 +29,7 @@ class _Response: self.headers = dict(headers or {}) self.url = url self.chunk_size = chunk_size + self.closed = False def iter_content(self, chunk_size=65536): if self.chunk_size: @@ -37,6 +38,9 @@ class _Response: return yield self.content + def close(self): + self.closed = True + class ImageStudioImageTests(TempDirMixin, unittest.TestCase): def _png_bytes(self, size=(20, 16), color=(80, 120, 200, 255)): @@ -246,6 +250,33 @@ class ImageStudioImageTests(TempDirMixin, unittest.TestCase): self.assert_removed(temp_dir) + def test_download_remote_image_cancels_between_chunks_and_closes_response(self): + png = self._png_bytes(size=(120, 120)) + response = _Response( + content=png, + headers={"Content-Type": "image/png"}, + chunk_size=32, + ) + session = SimpleNamespace(get=lambda *args, **kwargs: response) + checks = {"count": 0} + + def should_stop(): + checks["count"] += 1 + return checks["count"] >= 4 + + with mock.patch( + "app.image_studio_images.socket.getaddrinfo", + return_value=self._public_dns(), + ): + with self.assertRaises(image_studio_images.ImageStudioImageCancelled): + image_studio_images.download_remote_image( + "https://cdn.example.com/a.png", + session=session, + should_stop=should_stop, + ) + + self.assertTrue(response.closed) + def test_import_original_files_copies_valid_images_deduplicates_and_limits(self): with self.make_temp_dir() as temp_dir: db_path = os.path.join(temp_dir, "cmshopee.db") diff --git a/tests/test_product_suite_gui.py b/tests/test_product_suite_gui.py index 3ce0268..6901bc4 100644 --- a/tests/test_product_suite_gui.py +++ b/tests/test_product_suite_gui.py @@ -1866,6 +1866,298 @@ class ProductSuiteGuiTests(TempDirMixin, unittest.TestCase): self.assert_removed(temp_dir) + def test_pull_button_stops_with_confirmation_and_repeated_click_is_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) + tab = ProductSuiteTab(config=config, db_path=config["db_path"]) + self.addCleanup(tab.close) + state = tab._displayed_state + tab.item_id_edit.setText("51100639510") + + with mock.patch.object(tab, "_confirm", return_value=True), \ + mock.patch.object(tab, "_start_thread", return_value=mock.Mock()): + tab.pull_main_images() + + self.assertTrue(state.pull_running()) + self.assertEqual("停止拉取蝦皮", tab.pull_button.text()) + self.assertTrue(tab.pull_button.isEnabled()) + worker = state.pull_worker + + with mock.patch.object(tab, "_pull_stop_action", return_value="continue"): + tab.pull_main_images() + self.assertFalse(worker.is_cancelled()) + self.assertFalse(state.pull_stop_requested) + + with mock.patch.object(tab, "_pull_stop_action", return_value="keep"): + tab.pull_main_images() + self.assertTrue(worker.is_cancelled()) + self.assertTrue(state.pull_stop_requested) + self.assertEqual("正在停止...", tab.pull_button.text()) + + stop_action = mock.Mock(return_value="clear_current") + with mock.patch.object(tab, "_pull_stop_action", stop_action): + tab.pull_main_images() + stop_action.assert_not_called() + + state.pull_worker = None + state.pull_thread = None + tab._pull_run_states.clear() + state.pull_run_token = "" + + self.assert_removed(temp_dir) + + def test_stopping_pull_cancels_only_current_pull_downloads(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 + pull_worker = mock.Mock() + manual_worker = mock.Mock() + state.download_queue = [11, 12] + state.downloads = { + 13: (pull_worker, mock.Mock()), + 14: (manual_worker, mock.Mock()), + } + state.download_tokens = { + 11: "pull-token", + 12: "", + 13: "pull-token", + 14: "", + } + state.pull_download_asset_ids = {11, 13} + + tab._cancel_pull_downloads(state) + + self.assertEqual([12], state.download_queue) + self.assertEqual({13}, state.pull_download_asset_ids) + self.assertNotIn(11, state.download_tokens) + pull_worker.cancel.assert_called_once() + manual_worker.cancel.assert_not_called() + + self.assert_removed(temp_dir) + + def test_clear_stopped_pull_removes_only_new_unreferenced_remote_assets(self): + with self.make_temp_dir() as temp_dir: + config = self._config(temp_dir) + account = accounts.create_account( + "主店", + "alias-a", + debug_port=9222, + config=config, + ) + project = image_studio.create_or_get_project( + account, + item_id="51100639510", + path=config["db_path"], + ) + existing_remote = image_studio.sync_original_asset_urls( + project.id, + [{"index": 1, "src": "https://susercontent.com/existing.jpg"}], + path=config["db_path"], + )[0] + local_path = os.path.join(temp_dir, "local.png") + self._write_image(local_path) + local_asset = image_studio.add_asset( + project.id, + image_studio.ASSET_KIND_ORIGINAL, + local_path=local_path, + source_order=2, + path=config["db_path"], + ) + before_assets = image_studio.list_assets( + project.id, + kind=image_studio.ASSET_KIND_ORIGINAL, + path=config["db_path"], + ) + synced = image_studio.sync_original_asset_urls( + project.id, + [ + {"index": 1, "src": "https://susercontent.com/new-one.jpg"}, + {"index": 2, "src": "https://susercontent.com/new-two.jpg"}, + ], + path=config["db_path"], + ) + new_assets = [ + asset + for asset in synced + if asset.remote_url + and asset.remote_url.endswith(("new-one.jpg", "new-two.jpg")) + ] + referenced = next( + asset for asset in new_assets if asset.remote_url.endswith("new-two.jpg") + ) + image_studio.create_job( + project.id, + source_asset_id=referenced.id, + 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.pull_run_token = "clear-pull" + state.pull_stop_requested = True + state.pull_cleanup_mode = "clear_current" + state.pull_before_asset_ids = {int(asset.id) for asset in before_assets} + state.pull_before_asset_states = [ + { + "id": int(asset.id), + "status": asset.status, + "source_order": int(asset.source_order), + } + for asset in before_assets + ] + state.pull_asset_ids = {int(asset.id) for asset in new_assets} + state.pull_started_at = time.monotonic() + tab._pull_run_states["clear-pull"] = state.key + messages = [] + tab._message = lambda title, message, **kwargs: messages.append( + (title, message) + ) + + self.assertTrue(tab._finalize_pull(state, "clear-pull")) + + remaining = { + asset.id: asset + for asset in image_studio.list_assets( + project.id, + kind=image_studio.ASSET_KIND_ORIGINAL, + path=config["db_path"], + ) + } + removed = next( + asset for asset in new_assets if asset.id != referenced.id + ) + self.assertNotIn(removed.id, remaining) + self.assertIn(referenced.id, remaining) + self.assertEqual( + image_studio.ASSET_STATUS_AVAILABLE, + remaining[existing_remote.id].status, + ) + self.assertEqual(1, remaining[existing_remote.id].source_order) + self.assertIn(local_asset.id, remaining) + self.assertGreater( + remaining[referenced.id].source_order, + remaining[local_asset.id].source_order, + ) + self.assertEqual("拉取蝦皮主图已停止", messages[-1][0]) + self.assertIn("清理1张", messages[-1][1]) + self.assertIn("因引用保留1张", messages[-1][1]) + + self.assert_removed(temp_dir) + + def test_old_pull_result_does_not_override_current_run(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.pull_run_token = "current-pull" + state.pull_worker = mock.Mock() + tab._pull_run_states["current-pull"] = state.key + + tab._on_pull_finished( + "old-pull", + {"project": None, "assets": [], "count": 0}, + ) + + self.assertEqual("current-pull", state.pull_run_token) + self.assertIsNotNone(state.pull_worker) + + state.pull_worker = None + state.pull_run_token = "" + tab._pull_run_states.clear() + + self.assert_removed(temp_dir) + + def test_pull_thread_finished_fallback_finalizes_requested_stop(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.pull_run_token = "pull-fallback" + state.pull_stop_requested = True + state.pull_worker = mock.Mock() + state.pull_thread = mock.Mock() + state.pull_started_at = time.monotonic() + tab._pull_run_states["pull-fallback"] = state.key + messages = [] + tab._message = lambda title, message, **kwargs: messages.append( + (title, message) + ) + + tab._handle_pull_thread_finished("pull-fallback") + + self.assertEqual("", state.pull_run_token) + self.assertIsNone(state.pull_worker) + self.assertEqual("拉取蝦皮主图已停止", messages[-1][0]) + self.assertEqual("拉取蝦皮主图", tab.pull_button.text()) + + self.assert_removed(temp_dir) + + def test_pull_real_qthread_completion_restores_button_once(self): + with self.make_temp_dir() as temp_dir: + config = self._config(temp_dir) + account = accounts.create_account( + "主店", + "alias-a", + debug_port=9222, + config=config, + ) + project = image_studio.create_or_get_project( + account, + item_id="51100639510", + path=config["db_path"], + ) + tab = ProductSuiteTab(config=config, db_path=config["db_path"]) + self.addCleanup(tab.close) + state = tab._displayed_state + tab.item_id_edit.setText(project.item_id) + + with mock.patch.object(tab, "_confirm", return_value=True), \ + mock.patch( + "app.gui.workers.image_studio.pull_remote_main_image_urls", + return_value={ + "project": project, + "assets": [], + "account": account, + }, + ): + tab.pull_main_images() + pull_thread = state.pull_thread + deadline = time.monotonic() + 3 + while state.pull_running() and time.monotonic() < deadline: + QTest.qWait(20) + self.app.processEvents() + while pull_thread is not None and time.monotonic() < deadline: + try: + running = pull_thread.isRunning() + except RuntimeError: + pull_thread = None + break + if not running: + break + QTest.qWait(20) + self.app.processEvents() + + self.assertFalse(state.pull_running()) + self.assertIsNone(state.pull_worker) + self.assertEqual("拉取蝦皮主图", tab.pull_button.text()) + if pull_thread is not None: + self.assertFalse(pull_thread.isRunning()) + + self.assert_removed(temp_dir) + def test_generation_confirmation_prevents_job_creation_and_retry_bypasses_it(self): with self.make_temp_dir() as temp_dir: config = self._config(temp_dir) diff --git a/tests/test_workers.py b/tests/test_workers.py index 0434c60..27b3a2a 100644 --- a/tests/test_workers.py +++ b/tests/test_workers.py @@ -7,7 +7,7 @@ os.environ.setdefault("QT_QPA_PLATFORM", "offscreen") from _helpers import REPO_ROOT # noqa: F401 -from app import workers +from app import image_studio_images, workers if workers.QT_IMPORT_ERROR is not None: raise unittest.SkipTest("PySide6 未安装") @@ -16,7 +16,10 @@ from PySide6.QtCore import QEventLoop, QTimer from PySide6.QtWidgets import QApplication from app.workers import BaseWorker, run_worker -from app.gui.workers import ImageStudioDownloadOriginalWorker +from app.gui.workers import ( + ImageStudioDownloadOriginalWorker, + ImageStudioPullImagesWorker, +) class DemoWorker(BaseWorker): @@ -165,6 +168,34 @@ class WorkerTests(unittest.TestCase): self.assertEqual("蝦皮原主图下载失败，请稍后再次点击图片重试。", summary["error"]) self.assertNotIn("https://", summary["error"]) + def test_image_studio_pull_worker_returns_token_when_cancelled(self): + worker = ImageStudioPullImagesWorker( + "alias-a", + "51100639510", + pull_run_token="pull-token", + ) + worker.cancel() + + with mock.patch( + "app.gui.workers.image_studio.pull_remote_main_image_urls", + ) as pull: + summary = worker.execute() + + pull.assert_not_called() + self.assertTrue(summary["cancelled"]) + self.assertEqual("pull-token", summary["pull_run_token"]) + + def test_image_studio_original_download_maps_safe_cancel(self): + worker = ImageStudioDownloadOriginalWorker(12) + + with mock.patch( + "app.gui.workers.image_studio_images.download_original_asset", + side_effect=image_studio_images.ImageStudioImageCancelled("停止"), + ): + summary = worker.execute() + + self.assertEqual({"asset_id": 12, "cancelled": True}, summary) + if __name__ == "__main__": unittest.main()