fix(product-suite): finalize and cancel generation

app/image_studio_generation.py

@@ -6,7 +6,7 @@ import os
 import threading
 import time
 import urllib.parse
-from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
+from concurrent.futures import CancelledError, FIRST_COMPLETED, ThreadPoolExecutor, wait
 
 from . import ai, appconfig, image_studio
 from .version import APP_VERSION
@@ -185,7 +185,11 @@ def run_jobs(
             for job in job_list
         }
         while futures:
-            done, _ = wait(set(futures), return_when=FIRST_COMPLETED)
+            done, _ = wait(
+                set(futures),
+                timeout=0.2,
+                return_when=FIRST_COMPLETED,
+            )
             for future in done:
                 futures.pop(future)
                 try:
@@ -252,7 +256,14 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
         request_result = _poll_job(job.id, request_result["task_id"], runtime, request_result, db_path, should_stop, on_event)
         _raise_if_stopped(should_stop)
         out_path = _output_path(project, job, image_root)
-        saved_path = _download_and_save_job_image(request_result, out_path, config, on_event, job.id)
+        saved_path = _download_and_save_job_image(
+            request_result,
+            out_path,
+            config,
+            on_event,
+            job.id,
+            should_stop=should_stop,
+        )
         try:
             _raise_if_stopped(should_stop)
         except ImageStudioGenerationError:
@@ -281,17 +292,27 @@ def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, sho
         _notify(on_event, {"job_id": job.id, "step": "job_done", "result": "success"})
         return {"job": updated, "asset": asset, "status": "succeeded"}
     except Exception as exc:
-        status = "cancelled" if "停止" in str(exc) else "failed"
+        cancelled = isinstance(exc, CancelledError) or "停止" in str(exc)
+        status = "cancelled" if cancelled else "failed"
+        error = "用户已停止，已提交任务可稍后继续查询" if cancelled else str(exc)
         current_job = image_studio.get_job(job.id, path=db_path)
         updated = image_studio.update_job_status(
             job.id,
             status,
-            error=str(exc),
+            error=error,
             recovery_action=_recovery_action_for_job(current_job),
             path=db_path,
         )
-        _notify(on_event, {"job_id": job.id, "step": "job_done", "result": status, "detail": str(exc)})
-        return {"job": updated, "status": status, "error": str(exc)}
+        _notify(
+            on_event,
+            {
+                "job_id": job.id,
+                "step": "job_done",
+                "result": status,
+                "detail": error,
+            },
+        )
+        return {"job": updated, "status": status, "error": error}
 
 
 def _submit_or_resume_job(
@@ -395,6 +416,7 @@ def _poll_job(job_id, task_id, runtime, request_result, db_path, should_stop, on
             read_timeout=ai.CMHUB_IMAGE_POLL_READ_TIMEOUT_SECONDS,
             headers_extra={"X-Client-Version": str(APP_VERSION)},
         )
+        _raise_if_stopped(should_stop)
         status = str(data.get("status") or "").strip().lower()
         if status in {"queued", "running"}:
             _notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": status, "task_id": task_id})
@@ -424,21 +446,44 @@ def _poll_job(job_id, task_id, runtime, request_result, db_path, should_stop, on
         raise ImageStudioGenerationError("cmhub 生图任务状态返回格式错误")
 
 
-def _download_and_save_job_image(request_result, out_path, config, on_event, job_id):
+def _download_and_save_job_image(
+    request_result,
+    out_path,
+    config,
+    on_event,
+    job_id,
+    should_stop=None,
+):
+    _raise_if_stopped(should_stop)
     _notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "start"})
-    image_bytes, _ = ai._download_cmhub_image_with_retry(
-        request_result["image_url"],
-        connect_timeout=request_result["connect_timeout"],
-        read_timeout=request_result["read_timeout"],
-        use_system_proxy=request_result.get("use_system_proxy", False),
-        download_with_curl=request_result.get("download_with_curl", "auto"),
-    )
+    try:
+        image_bytes, _ = ai._download_cmhub_image_with_retry(
+            request_result["image_url"],
+            connect_timeout=request_result["connect_timeout"],
+            read_timeout=request_result["read_timeout"],
+            use_system_proxy=request_result.get("use_system_proxy", False),
+            download_with_curl=request_result.get("download_with_curl", "auto"),
+            should_stop=should_stop,
+        )
+    except CancelledError as exc:
+        raise ImageStudioGenerationError(
+            "用户已停止，已提交任务可稍后继续查询"
+        ) from exc
+    _raise_if_stopped(should_stop)
     saved_path = ai._save_jpeg(
         image_bytes,
         out_path,
         request_result.get("resolution") or appconfig.ai_config(config).get("resolution", "1k"),
         request_result.get("quality") or appconfig.ai_config(config).get("jpg_quality", 90),
     )
+    try:
+        _raise_if_stopped(should_stop)
+    except ImageStudioGenerationError:
+        try:
+            if os.path.isfile(saved_path):
+                os.remove(saved_path)
+        finally:
+            raise
     _notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "success"})
     return saved_path