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
@@ -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: