feat(product-suite): make image pulls cancellable

2026-07-16 19:37:22 +08:00
parent c9afaad24b
commit 75af87d441
12 changed files with 1157 additions and 88 deletions
@@ -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)
@@ -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: