fix: 支持删除图片异常批次

2026-07-15 17:52:29 +08:00
parent ff0a93677a
commit 755e210951
9 changed files with 531 additions and 51 deletions
@@ -1732,6 +1732,122 @@ class WorkbenchAndUnreadTests(TestCase):
        self.assertEqual(body["total"], 0)
        self.assertEqual(body["by_product"], {})

+    def test_exception_batch_delete_restore_and_purge_are_atomic(self):
+        """部分失败批次进入“图片异常批次”，而非与资产垃圾桶重复；恢复后整批回放。"""
+        tasks = self._generate(mode="cover", count=2, platform_id="douyin")
+        failed, succeeded = tasks
+        # 模拟其中一张供应商失败：失败任务无成图，另一张保持成功，构成真实的混合批次。
+        Asset.objects.filter(origin_task=failed).delete()
+        failed.status = AITask.Status.FAILED
+        failed.save(update_fields=["status"])
+        succeeded_asset = Asset.objects.get(origin_task=succeeded)
+
+        deleted = self.client.delete(f"/api/ai/tasks/{failed.id}/workbench-batch/")
+        self.assertEqual(deleted.status_code, 200, deleted.content)
+        self.assertEqual(deleted.json()["storage"], "exception_batch")
+        failed.refresh_from_db()
+        succeeded.refresh_from_db()
+        succeeded_asset.refresh_from_db()
+        self.assertTrue(failed.is_deleted)
+        self.assertTrue(succeeded.is_deleted)
+        self.assertTrue(succeeded_asset.is_deleted)
+        self.assertEqual(
+            self.client.get(f"/api/ai/tasks/workbench/?mode=cover&product={self.product.id}").json()["tasks"], []
+        )
+
+        trash = self.client.get("/api/ai/tasks/workbench-exception-batches-trash/")
+        self.assertEqual(trash.status_code, 200, trash.content)
+        rows = trash.json()["results"]
+        self.assertEqual(len(rows), 1)
+        self.assertEqual(rows[0]["id"], str(failed.id))
+        self.assertEqual(rows[0]["cover_preview_url"], "http://x/cover.png")
+        asset_trash_ids = {row["id"] for row in self.client.get("/api/assets/trash/").json()["results"]}
+        self.assertNotIn(str(succeeded_asset.id), asset_trash_ids)
+
+        restored = self.client.post(f"/api/ai/tasks/{failed.id}/restore-workbench-batch/")
+        self.assertEqual(restored.status_code, 200, restored.content)
+        failed.refresh_from_db()
+        succeeded.refresh_from_db()
+        succeeded_asset.refresh_from_db()
+        self.assertFalse(failed.is_deleted)
+        self.assertFalse(succeeded.is_deleted)
+        self.assertFalse(succeeded_asset.is_deleted)
+        rows = self.client.get(f"/api/ai/tasks/workbench/?mode=cover&product={self.product.id}").json()["tasks"]
+        self.assertEqual({row["id"] for row in rows}, {str(failed.id), str(succeeded.id)})
+
+        self.client.delete(f"/api/ai/tasks/{failed.id}/workbench-batch/")
+        purged = self.client.delete(f"/api/ai/tasks/{failed.id}/purge-workbench-batch/")
+        self.assertEqual(purged.status_code, 204, purged.content)
+        failed.refresh_from_db()
+        succeeded.refresh_from_db()
+        succeeded_asset.refresh_from_db()
+        self.assertIsNotNone(failed.purged_at)
+        self.assertIsNotNone(succeeded.purged_at)
+        self.assertIsNotNone(succeeded_asset.purged_at)
+        self.assertEqual(self.client.get("/api/ai/tasks/workbench-exception-batches-trash/").json()["results"], [])
+
+    def test_successful_batch_keeps_existing_asset_trash_behavior(self):
+        tasks = self._generate(mode="cover", count=2, platform_id="douyin")
+        assets = list(Asset.objects.filter(origin_task__in=tasks))
+
+        deleted = self.client.delete(f"/api/ai/tasks/{tasks[0].id}/workbench-batch/")
+        self.assertEqual(deleted.status_code, 200, deleted.content)
+        self.assertEqual(deleted.json()["storage"], "asset")
+        self.assertFalse(AITask.objects.get(id=tasks[0].id).is_deleted)
+        self.assertTrue(all(Asset.objects.get(id=asset.id).is_deleted for asset in assets))
+        asset_trash_ids = {row["id"] for row in self.client.get("/api/assets/trash/").json()["results"]}
+        self.assertTrue({str(asset.id) for asset in assets}.issubset(asset_trash_ids))
+        self.assertEqual(self.client.get("/api/ai/tasks/workbench-exception-batches-trash/").json()["results"], [])
+
+    def test_image_conversation_exception_batch_hides_and_restores_only_its_tasks(self):
+        """图片创作走对话任务读取，也必须在异常批次删除/恢复后保持一致。"""
+        conversation = ImageConversation.objects.create(
+            team=self.team, created_by=self.user, mode=ImageConversation.Mode.IMAGE, title="失败图片创作", product=self.product
+        )
+        task = AITask.objects.create(
+            team=self.team,
+            created_by=self.user,
+            conversation=conversation,
+            task_type=AITask.Type.PRODUCT_IMAGE,
+            status=AITask.Status.FAILED,
+            model_config=ModelConfig.objects.filter(capability=ModelConfig.Capability.IMAGE).first(),
+            idempotency_key="conversation-exception-batch",
+            request_payload={"mode": "image", "batch_id": "conversation-exception-batch", "prompt": "失败的图片创作", "product_id": str(self.product.id)},
+        )
+        task_rows = self.client.get(f"/api/ai/image-conversations/{conversation.id}/tasks/").json()["tasks"]
+        self.assertEqual([row["id"] for row in task_rows], [str(task.id)])
+
+        self.assertEqual(self.client.delete(f"/api/ai/tasks/{task.id}/workbench-batch/").status_code, 200)
+        self.assertEqual(self.client.get(f"/api/ai/image-conversations/{conversation.id}/tasks/").json()["tasks"], [])
+
+        self.assertEqual(self.client.post(f"/api/ai/tasks/{task.id}/restore-workbench-batch/").status_code, 200)
+        task_rows = self.client.get(f"/api/ai/image-conversations/{conversation.id}/tasks/").json()["tasks"]
+        self.assertEqual([row["id"] for row in task_rows], [str(task.id)])
+
+    def test_inflight_batch_cannot_be_deleted(self):
+        tasks = self._generate(mode="cover", count=1, platform_id="douyin")
+        task = tasks[0]
+        task.status = AITask.Status.POLLING
+        task.save(update_fields=["status"])
+
+        deleted = self.client.delete(f"/api/ai/tasks/{task.id}/workbench-batch/")
+        self.assertEqual(deleted.status_code, 409, deleted.content)
+        task.refresh_from_db()
+        self.assertFalse(task.is_deleted)
+
+    def test_other_team_cannot_delete_image_batch(self):
+        task = self._generate(mode="cover", count=1, platform_id="douyin")[0]
+        other = User.objects.create_user(username="other-batch-owner", password="pass")
+        other_team = Team.objects.create(name="Other batch team", owner=other)
+        TeamMember.objects.create(team=other_team, user=other, role="owner", status="active")
+        other_client = APIClient()
+        other_client.force_authenticate(other)
+
+        deleted = other_client.delete(f"/api/ai/tasks/{task.id}/workbench-batch/")
+        self.assertEqual(deleted.status_code, 404, deleted.content)
+        task.refresh_from_db()
+        self.assertFalse(task.is_deleted)
+
    def test_unread_is_team_scoped_and_mode_whitelisted(self):
        self._generate(mode="cover", count=1)
        # 非生成任务(脚本 agent 复用 mode=auto)不计入未读
@@ -1,5 +1,5 @@
 from django.db import transaction
-from django.db.models import Count, Exists, OuterRef
+from django.db.models import Count, Exists, OuterRef, Q
 from django.utils import timezone
 from rest_framework import status
 from rest_framework.decorators import action
@@ -9,9 +9,10 @@ from rest_framework.views import APIView
 from rest_framework.viewsets import ModelViewSet, ReadOnlyModelViewSet

 from apps.assets.models import Asset
-from apps.assets.serializers import AssetSerializer
+from apps.assets.serializers import AssetFileSerializer, AssetSerializer
 from apps.common.api import TeamScopedViewSetMixin, get_current_team
 from apps.common.celery_health import require_worker
+from apps.products.models import Product

 from .generation_errors import classify_generation_error, public_error_for_task
 from .models import AITask, ImageConversation, ModelConfig
@@ -146,6 +147,144 @@ class AITaskViewSet(TeamScopedViewSetMixin, ReadOnlyModelViewSet):
    # YYX#row22:只有「工作台图片生成」(mode∈model/cover/image)才计入未读;
    # 脚本/实体抽取/故事板等流水线内部任务不进这块未读统计。
    _GEN_MODES = ("model", "cover", "image")
+    _WORKBENCH_IMAGE_TYPES = (AITask.Type.PERSON_IMAGE, AITask.Type.PRODUCT_IMAGE)
+    _DELETABLE_IMAGE_BATCH_STATUSES = (AITask.Status.SUCCEEDED, AITask.Status.FAILED, AITask.Status.CANCELLED)
+
+    def _image_batch_tasks(self, pk, *, deleted: bool):
+        """Resolve one standalone image batch from a user-visible task anchor.
+
+        The browser only supplies an anchor task id. Team, mode and batch membership are always
+        derived server-side so one team cannot delete or restore another team's task group.
+        """
+        base = AITask.objects.filter(
+            team=self.get_team(),
+            project__isnull=True,
+            task_type__in=self._WORKBENCH_IMAGE_TYPES,
+            is_deleted=deleted,
+            purged_at__isnull=True,
+        )
+        anchor = base.filter(pk=pk).first()
+        if anchor is None:
+            return None, []
+        payload = anchor.request_payload or {}
+        mode = str(payload.get("mode") or "").strip()
+        if mode not in self._GEN_MODES:
+            return None, []
+        batch_id = str(payload.get("batch_id") or "").strip()
+        batch_qs = base.filter(request_payload__mode=mode)
+        if batch_id:
+            batch_qs = batch_qs.filter(request_payload__batch_id=batch_id)
+        else:
+            batch_qs = batch_qs.filter(pk=anchor.pk)
+        return anchor, list(batch_qs.order_by("created_at").select_for_update())
+
+    @staticmethod
+    def _is_exception_batch(tasks):
+        return any(task.status in (AITask.Status.FAILED, AITask.Status.CANCELLED) for task in tasks)
+
+    @action(detail=True, methods=["delete"], url_path="workbench-batch")
+    def delete_workbench_batch(self, request, pk=None):
+        """Delete a completed image batch.
+
+        Fully successful batches preserve the existing asset-trash behavior. A batch containing
+        a failed/cancelled task is soft-deleted as one recoverable exception batch instead.
+        """
+        with transaction.atomic():
+            anchor, tasks = self._image_batch_tasks(pk, deleted=False)
+            if anchor is None or not tasks:
+                return Response({"detail": "图片批次不存在或不可删除"}, status=status.HTTP_404_NOT_FOUND)
+            if any(task.status not in self._DELETABLE_IMAGE_BATCH_STATUSES for task in tasks):
+                return Response({"detail": "批次仍在处理中，暂不能删除"}, status=status.HTTP_409_CONFLICT)
+
+            now = timezone.now()
+            task_ids = [task.id for task in tasks]
+            Asset.objects.filter(origin_task_id__in=task_ids, purged_at__isnull=True).update(is_deleted=True, updated_at=now)
+            if self._is_exception_batch(tasks):
+                AITask.objects.filter(id__in=task_ids).update(is_deleted=True, updated_at=now)
+                return Response({"storage": "exception_batch", "deleted_task_count": len(task_ids)})
+            return Response({"storage": "asset", "deleted_task_count": 0})
+
+    @action(detail=True, methods=["post"], url_path="restore-workbench-batch")
+    def restore_workbench_batch(self, request, pk=None):
+        """Restore one previously deleted exception batch, including its generated images."""
+        with transaction.atomic():
+            anchor, tasks = self._image_batch_tasks(pk, deleted=True)
+            if anchor is None or not tasks:
+                return Response({"detail": "图片异常批次不存在或无法恢复"}, status=status.HTTP_404_NOT_FOUND)
+            now = timezone.now()
+            task_ids = [task.id for task in tasks]
+            AITask.objects.filter(id__in=task_ids).update(is_deleted=False, purged_at=None, updated_at=now)
+            Asset.objects.filter(origin_task_id__in=task_ids, purged_at__isnull=True).update(is_deleted=False, updated_at=now)
+        return Response({"restored_task_count": len(task_ids)})
+
+    @action(detail=True, methods=["delete"], url_path="purge-workbench-batch")
+    def purge_workbench_batch(self, request, pk=None):
+        """Permanently hide one deleted exception batch and all of its generated images."""
+        with transaction.atomic():
+            anchor, tasks = self._image_batch_tasks(pk, deleted=True)
+            if anchor is None or not tasks:
+                return Response({"detail": "图片异常批次不存在或已彻底删除"}, status=status.HTTP_404_NOT_FOUND)
+            now = timezone.now()
+            task_ids = [task.id for task in tasks]
+            AITask.objects.filter(id__in=task_ids).update(is_deleted=True, purged_at=now, updated_at=now)
+            Asset.objects.filter(origin_task_id__in=task_ids, purged_at__isnull=True).update(is_deleted=True, purged_at=now, updated_at=now)
+        return Response(status=status.HTTP_204_NO_CONTENT)
+
+    @action(detail=False, methods=["get"], url_path="workbench-exception-batches-trash")
+    def workbench_exception_batches_trash(self, request):
+        """List recoverable failed/cancelled image batches for the existing global trash page."""
+        tasks = list(
+            AITask.objects.filter(
+                team=self.get_team(),
+                project__isnull=True,
+                task_type__in=self._WORKBENCH_IMAGE_TYPES,
+                is_deleted=True,
+                purged_at__isnull=True,
+                request_payload__mode__in=self._GEN_MODES,
+            ).filter(Q(conversation__isnull=True) | Q(conversation__is_deleted=False)).order_by("created_at")
+        )
+        groups = {}
+        product_ids = set()
+        for task in tasks:
+            payload = task.request_payload or {}
+            key = str(payload.get("batch_id") or task.id)
+            groups.setdefault(key, []).append(task)
+            product_id = str(payload.get("product_id") or "").strip()
+            if product_id:
+                product_ids.add(product_id)
+
+        products = Product.objects.filter(
+            team=self.get_team(), id__in=product_ids, status=Product.Status.ACTIVE, purged_at__isnull=True
+        ).select_related("cover_asset").prefetch_related("cover_asset__files")
+        product_by_id = {str(product.id): product for product in products}
+
+        rows = []
+        for key, group in groups.items():
+            if not self._is_exception_batch(group):
+                continue
+            first = group[0]
+            payload = first.request_payload or {}
+            product = product_by_id.get(str(payload.get("product_id") or ""))
+            cover = ""
+            if product and product.cover_asset and not product.cover_asset.is_deleted and product.cover_asset.purged_at is None:
+                files = list(product.cover_asset.files.all())
+                primary = next((item for item in files if item.is_primary), files[0] if files else None)
+                if primary:
+                    cover = AssetFileSerializer(primary).data.get("preview_url", "")
+            intended = sum(1 for task in group if not (task.request_payload or {}).get("batch_append"))
+            rows.append({
+                "id": str(first.id),
+                "batch_id": str(payload.get("batch_id") or ""),
+                "mode": str(payload.get("mode") or "image"),
+                "prompt": str(payload.get("prompt") or ""),
+                "count": max(1, intended),
+                "product_id": str(payload.get("product_id") or ""),
+                "product_title": product.title if product else "",
+                "cover_preview_url": cover,
+                "updated_at": max(task.updated_at for task in group),
+            })
+        rows.sort(key=lambda row: row["updated_at"], reverse=True)
+        return Response({"results": rows})

    def _unread_base(self):
        """本团队、属于图片生成、且未读(read_at is null)的任务集合(已 annotate rp_mode/rp_product_id)。
@@ -187,6 +326,8 @@ class AITaskViewSet(TeamScopedViewSetMixin, ReadOnlyModelViewSet):
                team=self.get_team(),
                project__isnull=True,  # 工作台独立生图不挂项目;排除流水线内部任务
                task_type__in=[AITask.Type.PERSON_IMAGE, AITask.Type.PRODUCT_IMAGE],
+                is_deleted=False,
+                purged_at__isnull=True,
            )
            .defer("request_payload", "response_payload")
            .annotate(
@@ -326,7 +467,7 @@ class ImageConversationViewSet(TeamScopedViewSetMixin, ModelViewSet):

        conversation = self.get_object()
        tasks = (
-            AITask.objects.filter(conversation=conversation)
+            AITask.objects.filter(conversation=conversation, is_deleted=False, purged_at__isnull=True)
            .prefetch_related("generated_assets", "generated_assets__files")
            .order_by("created_at")
        )