feat: 优化图片资源删除管理，生图任务列表管理

server/services/generator.py

@@ -13,6 +13,7 @@ import base64
 import logging
 import mimetypes
 import re
+from dataclasses import dataclass
 from typing import Awaitable, Callable
 
 import httpx
@@ -30,7 +31,7 @@ from services.tasks import (
     TaskImage,
 )
 from services.watermark import apply_watermark
-from services.storage import download_bytes, get_storage, local_path
+from services.storage import download_bytes, get_storage, local_path, read_local_bytes
 
 log = logging.getLogger("suite.generator")
 
@@ -125,7 +126,7 @@ async def _resolve_ref_bytes(url: str) -> tuple[bytes, str]:
     path = local_path(url)
     if path is not None:
         mime = mimetypes.guess_type(path.name)[0] or "image/jpeg"
-        return path.read_bytes(), mime
+        return read_local_bytes(path), mime
     if url.startswith(("http://", "https://")):
         from api.proxy import guess_referer
         data, ctype = await download_bytes(url, referer=guess_referer(url))
@@ -288,6 +289,18 @@ async def _rightapi_poll_task(client: httpx.AsyncClient, headers: dict, origin:
     raise TimeoutError(f"RightAPI 异步任务超时 ({max_wait}s): task_id={task_id}")
 
 
+async def _rightapi_poll_task_once(client: httpx.AsyncClient, headers: dict, origin: str, task_id: str) -> dict | None:
+    """持久化 Worker 的单次查询版本；未完成返回 None，不占住 Worker 等待。"""
+    response = await client.get(f"{origin}/v1/tasks/{task_id}", headers=headers, timeout=30)
+    _raise_api_error(response, "RightAPI")
+    result = response.json()
+    status = result.get("status", "")
+    if status in ("failed", "error", "cancelled"):
+        error = result.get("error") or {}
+        raise RuntimeError(f"RightAPI 任务失败: {error.get('message') or result}")
+    return result if "data" in result else None
+
+
 def _rightapi_extract_image(result: dict) -> tuple[str | None, str | None]:
     """从轮询完成结果里取 (kind, payload)：kind ∈ url | b64，未取到返回 (None, None)。
 
@@ -411,6 +424,19 @@ async def _pcore_poll_task(client: httpx.AsyncClient, headers: dict, base: str,
     raise TimeoutError(f"Pcore 异步任务超时 ({max_wait}s): task_id={task_id}")
 
 
+async def _pcore_poll_task_once(client: httpx.AsyncClient, headers: dict, base: str,
+                                kind: str, task_id: str) -> dict | None:
+    """持久化 Worker 的单次查询版本；未完成返回 None。"""
+    response = await client.get(f"{base}/v1/images/{kind}/{task_id}", headers=headers, timeout=30)
+    _raise_api_error(response, "Pcore")
+    result = response.json()
+    status = result.get("status", "")
+    if status in ("failed", "error", "cancelled"):
+        error = result.get("error") or {}
+        raise RuntimeError(f"Pcore 任务失败: {error.get('message') or result}")
+    return result if status in ("completed", "succeeded") else None
+
+
 async def _pcore_request(s, prompt: str, ref_images: list[str], size: str, model: str) -> bytes:
     """Pcore：无参考图走 /v1/images/generations，带参考图走 /v1/images/edits（JSON，data-URI）。
 
@@ -482,6 +508,164 @@ async def generate_pcore(prompt: str, ref_images: list[str], size: str = "1536x2
     raise last_exc  # type: ignore[misc]
 
 
+# ── 可恢复的逐步 provider 调用（持久化 Worker 使用）──────────────────────
+
+@dataclass(frozen=True)
+class GenerationAdvance:
+    """一次 Worker 推进的结果：data 有值时已完成，否则持久化外部 task id 后稍后再查。"""
+
+    data: bytes | None = None
+    provider_task_id: str | None = None
+    provider_task_kind: str | None = None
+
+
+async def _download_provider_image(client: httpx.AsyncClient, result: dict, provider: str) -> bytes:
+    kind, payload = _rightapi_extract_image(result)
+    if kind == "b64" and payload:
+        return base64.b64decode(payload.split(",", 1)[-1] if "," in payload else payload)
+    if kind == "url" and payload:
+        response = await client.get(payload, timeout=get_settings().request_timeout)
+        _raise_api_error(response, provider)
+        return response.content
+    raise RuntimeError(f"{provider} 任务完成但没有图片数据: {result}")
+
+
+async def _poll_tongyi_once(client: httpx.AsyncClient, key: str, task_id: str) -> tuple[str, str | None]:
+    response = await client.get(
+        "https://dashscope.aliyuncs.com/api/v1/tasks/" + task_id,
+        headers={"Authorization": f"Bearer {key}"}, timeout=30,
+    )
+    _raise_api_error(response, "通义")
+    result = response.json()
+    status = result.get("output", {}).get("task_status", "")
+    if status in ("FAILED", "UNKNOWN"):
+        raise RuntimeError(f"通义任务失败: {result}")
+    if status != "SUCCEEDED":
+        return "waiting", None
+    choices = result.get("output", {}).get("choices", [])
+    if choices:
+        content = choices[0].get("message", {}).get("content", [])
+        if content:
+            return "done", content[0].get("image", "")
+    results = result.get("output", {}).get("results", [])
+    if results:
+        return "done", results[0].get("url") or results[0].get("b64_image", "")
+    raise RuntimeError(f"通义任务成功但无结果: {result}")
+
+
+async def advance_generation(
+    provider: str,
+    prompt: str,
+    ref_images: list[str],
+    *,
+    size: str,
+    model: str,
+    provider_task_id: str | None = None,
+    provider_task_kind: str | None = None,
+) -> GenerationAdvance:
+    """推进一次图片生成，不在本函数内无限轮询。
+
+    对异步供应商：首次仅提交并返回 task id；后续调用只轮询该 id。这样 Worker 重启后
+    能继续外部任务，且在等待期间可以领取其他任务。豆包和同步千问仍在一次调用内完成。
+    """
+    settings = get_settings()
+    if provider == "doubao":
+        return GenerationAdvance(data=await generate_doubao(prompt, ref_images, size=size, model=model))
+
+    if provider == "tongyi":
+        if not settings.dashscope_api_key:
+            raise RuntimeError("未配置 DASHSCOPE_API_KEY（.env）")
+        is_wan = _is_wan_model(model)
+        if not is_wan:
+            return GenerationAdvance(data=await generate_tongyi(prompt, ref_images, size=size, model=model))
+        headers = {"Authorization": f"Bearer {settings.dashscope_api_key}", "Content-Type": "application/json"}
+        if provider_task_id:
+            async with httpx.AsyncClient(timeout=60, verify=False) as client:
+                state, payload = await _poll_tongyi_once(client, settings.dashscope_api_key, provider_task_id)
+                if state == "waiting":
+                    return GenerationAdvance(provider_task_id=provider_task_id, provider_task_kind="tongyi_wan")
+                if payload and (payload.startswith("data:") or len(payload) > 500):
+                    return GenerationAdvance(data=base64.b64decode(payload.split(",", 1)[-1] if "," in payload else payload))
+                if not payload:
+                    raise RuntimeError("通义任务成功但图片地址为空")
+                response = await client.get(payload, timeout=settings.request_timeout)
+                _raise_api_error(response, "通义")
+                return GenerationAdvance(data=response.content)
+        url = settings.dashscope_base_url or "https://dashscope.aliyuncs.com/api/v1/services/aigc/image-generation/generation"
+        content: list[dict] = []
+        if not _is_t2i_model(model):
+            content = [{"image": await _resolve_ref(item)} for item in ref_images]
+        content.append({"text": prompt})
+        body = {
+            "model": model,
+            "input": {"messages": [{"role": "user", "content": content}]},
+            "parameters": {"size": size, "n": 1, "watermark": False, "negative_prompt": DEFAULT_NEGATIVE_PROMPT[:500]},
+        }
+        async with httpx.AsyncClient(timeout=settings.request_timeout, verify=False) as client:
+            response = await client.post(url, headers={**headers, "X-DashScope-Async": "enable"}, json=body)
+            _raise_api_error(response, "通义")
+            task_id = response.json().get("output", {}).get("task_id", "")
+            if not task_id:
+                raise RuntimeError(f"通义万象未返回 task_id: {response.json()}")
+            return GenerationAdvance(provider_task_id=task_id, provider_task_kind="tongyi_wan")
+
+    if provider == "rightapi":
+        if not settings.rightapi_api_key:
+            raise RuntimeError("未配置 RIGHTAPI_API_KEY（.env）")
+        base = settings.rightapi_base_url.rstrip("/")
+        origin = base.split("/draw", 1)[0].rstrip("/") or base
+        headers = {"Authorization": f"Bearer {settings.rightapi_api_key}"}
+        async with httpx.AsyncClient(timeout=max(settings.request_timeout, 600), verify=False) as client:
+            if provider_task_id:
+                result = await _rightapi_poll_task_once(client, headers, origin, provider_task_id)
+                if result is None:
+                    return GenerationAdvance(provider_task_id=provider_task_id, provider_task_kind="rightapi")
+                return GenerationAdvance(data=await _download_provider_image(client, result, "RightAPI"))
+            body: dict = {"model": model, "prompt": prompt, "n": 1, "size": size, "async": True}
+            if ref_images:
+                body["image"] = [await _resolve_ref(item) for item in ref_images]
+            response = await client.post(
+                f"{base}/v1/images/generations", headers={**headers, "Content-Type": "application/json"}, json=body,
+            )
+            _raise_api_error(response, "RightAPI")
+            submitted = response.json()
+            task_id = submitted.get("task_id") or ""
+            if task_id:
+                return GenerationAdvance(provider_task_id=task_id, provider_task_kind="rightapi")
+            return GenerationAdvance(data=await _download_provider_image(client, submitted, "RightAPI"))
+
+    if provider == "pcore":
+        if not settings.pcore_api_key:
+            raise RuntimeError("未配置 PCORE_API_KEY（.env）")
+        base = settings.pcore_base_url.rstrip("/")
+        headers = {"Authorization": f"Bearer {settings.pcore_api_key}"}
+        async with httpx.AsyncClient(timeout=max(settings.request_timeout, 600), verify=False) as client:
+            if provider_task_id:
+                kind = provider_task_kind or "edits"
+                result = await _pcore_poll_task_once(client, headers, base, kind, provider_task_id)
+                if result is None:
+                    return GenerationAdvance(provider_task_id=provider_task_id, provider_task_kind=kind)
+                return GenerationAdvance(data=await _download_provider_image(client, result, "Pcore"))
+            kind = "edits" if ref_images else "generations"
+            body: dict = {"model": model, "prompt": prompt, "n": 1, "size": size, "async": True}
+            if ref_images:
+                body["images"] = [{"image_url": await _resolve_ref(item)} for item in ref_images]
+            response = await client.post(
+                f"{base}/v1/images/{kind}", headers={**headers, "Content-Type": "application/json"}, json=body,
+            )
+            _raise_api_error(response, "Pcore")
+            submitted = response.json()
+            if isinstance(submitted, dict) and submitted.get("error"):
+                error = submitted["error"] or {}
+                raise RuntimeError(f"Pcore 提交失败: {error.get('message') or submitted}")
+            task_id = submitted.get("task_id") or submitted.get("id") or ""
+            if not task_id:
+                raise RuntimeError(f"Pcore 提交未返回任务 id: {submitted}")
+            return GenerationAdvance(provider_task_id=task_id, provider_task_kind=kind)
+
+    raise RuntimeError(f"未知 provider: {provider}")
+
+
 GENERATORS = {
     "doubao": generate_doubao,
     "tongyi": generate_tongyi,
@@ -524,12 +708,12 @@ def _refs_for_job(images: list[dict], job: dict) -> list[str]:
 _GEN_LOCK = asyncio.Lock()
 
 
-async def _store_image_bytes(task: Task, index: int, data: bytes) -> str:
+async def _store_image_bytes(task: Task, index: int, data: bytes, version: str = "result") -> str:
     """生成图字节 → 存储落盘，返回可访问 URL。按 PNG 魔数定扩展名（部分中转不遵守 output_format）。"""
     is_png = data[:8] == b"\x89PNG\r\n\x1a\n"
     ext = ".png" if is_png else ".jpg"
     ctype = "image/png" if is_png else "image/jpeg"
-    return await get_storage().save_bytes(data, f"suites/{task.id}/{index}{ext}", ctype)
+    return await get_storage().save_bytes(data, f"suites/{task.id}/{version}/{index}{ext}", ctype)
 
 
 async def run_suite(task: Task, on_image_ok: OnImageOk | None = None) -> None:
@@ -565,17 +749,24 @@ async def run_suite(task: Task, on_image_ok: OnImageOk | None = None) -> None:
             task.images.append(image)
             try:
                 # 提示词按模型家族分发：国产主体参考 / gpt edits 保真 / google 主体保持
+                # 单张 AI 编辑可传入 prompt_extra：任务展示名保留「AI 生图」，
+                # 但不把中文展示名拼进模型 prompt，避免生成图内出现中文。
+                prompt_extra = job.get("prompt_extra") or job
                 prompt = build_prompt(
                     provider_name, model, type_id, ctx, task.style_set, task.lang,
-                    extra=job, style_prompt=task.style_prompt, requirements=task.requirements,
+                    extra=prompt_extra, style_prompt=task.style_prompt, requirements=task.requirements,
                 )
                 refs = _refs_for_job(list(task.ref_images or []), job)
                 data = await generator(prompt, refs, size=size, model=model)
                 # 水印：AI 出图返回后、落盘前的后处理（失败不阻断，内部返回原图）
                 wm = task.watermark or {}
                 if wm.get("enabled"):
+                    image.source_url = await _store_image_bytes(task, len(task.images), data, version="original")
                     data = apply_watermark(data, wm)
-                image.url = await _store_image_bytes(task, len(task.images), data)
+                    image.url = await _store_image_bytes(task, len(task.images), data, version="watermarked")
+                    image.watermark_config = wm
+                else:
+                    image.url = await _store_image_bytes(task, len(task.images), data)
                 image.status = IMG_OK
                 ok += 1
                 if on_image_ok is not None and image.url: