fix(product-suite): finalize and cancel generation

app/ai.py

@@ -1447,6 +1447,7 @@ def _download_cmhub_image_with_retry(
     use_system_proxy=False,
     download_with_curl="false",
     on_step=None,
+    should_stop=None,
     attempts=CMHUB_IMAGE_DOWNLOAD_ATTEMPTS,
     slow_threshold=CMHUB_IMAGE_SLOW_DOWNLOAD_SECONDS,
 ):
@@ -1454,6 +1455,7 @@ def _download_cmhub_image_with_retry(
     total_started = time.perf_counter()
     last_exc = None
     for index in range(total_attempts):
+        _raise_if_download_cancelled(should_stop)
         try:
             image_bytes = _download_cmhub_image(
                 url,
@@ -1461,7 +1463,9 @@ def _download_cmhub_image_with_retry(
                 read_timeout=read_timeout,
                 use_system_proxy=use_system_proxy,
                 download_with_curl=download_with_curl,
+                should_stop=should_stop,
             )
+            _raise_if_download_cancelled(should_stop)
             elapsed = time.perf_counter() - total_started
             if elapsed >= float(slow_threshold or 0):
                 _notify_step_event(
@@ -1473,6 +1477,8 @@ def _download_cmhub_image_with_retry(
                     level="warning",
                 )
             return image_bytes, elapsed
+        except CancelledError:
+            raise
         except Exception as exc:
             last_exc = exc
             if index + 1 >= total_attempts or not _cmhub_download_retryable(exc):
@@ -1486,7 +1492,7 @@ def _download_cmhub_image_with_retry(
                 attempt=index + 1,
                 attempts=total_attempts,
             )
-            time.sleep(min(2.0, 0.5 * (index + 1)))
+            _sleep_download_retry(min(2.0, 0.5 * (index + 1)), should_stop)
     if total_attempts > 1 and _cmhub_download_retryable(last_exc):
         raise AIError(
             "下载 cmhub 图片失败（已尝试 %s 次）: %s"
@@ -1739,7 +1745,9 @@ def _download_cmhub_image(
     max_bytes=CMHUB_IMAGE_MAX_BYTES,
     use_system_proxy=False,
     download_with_curl="false",
+    should_stop=None,
 ):
+    _raise_if_download_cancelled(should_stop)
     _assert_public_http_url(url)
     if _should_use_curl_for_cmhub_download(download_with_curl):
         try:
@@ -1749,7 +1757,10 @@ def _download_cmhub_image(
                 read_timeout=read_timeout,
                 max_bytes=max_bytes,
                 use_system_proxy=use_system_proxy,
+                should_stop=should_stop,
             )
+        except CancelledError:
+            raise
         except AIError:
             pass
     return _download_cmhub_image_with_requests(
@@ -1757,32 +1768,53 @@ def _download_cmhub_image(
         connect_timeout=connect_timeout,
         read_timeout=read_timeout,
         max_bytes=max_bytes,
+        should_stop=should_stop,
     )
 
 
-def _download_cmhub_image_with_requests(url, connect_timeout, read_timeout, max_bytes):
+def _download_cmhub_image_with_requests(
+    url,
+    connect_timeout,
+    read_timeout,
+    max_bytes,
+    should_stop=None,
+):
+    response = None
     try:
+        _raise_if_download_cancelled(should_stop)
         response = _cmhub_session().get(
             url,
             stream=True,
             timeout=(max(1, int(connect_timeout)), max(1, int(read_timeout))),
         )
+        _raise_if_download_cancelled(should_stop)
+        status = getattr(response, "status_code", 200)
+        if status >= 400:
+            raise AIError("下载 cmhub 图片失败: HTTP %s" % status)
+        chunks = []
+        total = 0
+        iterator = (
+            response.iter_content(chunk_size=65536)
+            if hasattr(response, "iter_content")
+            else [response.content]
+        )
+        for chunk in iterator:
+            _raise_if_download_cancelled(should_stop)
+            if not chunk:
+                continue
+            total += len(chunk)
+            if total > max_bytes:
+                raise AIError("下载 cmhub 图片失败: 图片超过大小上限")
+            chunks.append(chunk)
+        _raise_if_download_cancelled(should_stop)
+        return b"".join(chunks)
+    except CancelledError:
+        raise
     except requests.exceptions.RequestException as exc:
         raise AIError("下载 cmhub 图片失败: %s" % exc) from exc
-    status = getattr(response, "status_code", 200)
-    if status >= 400:
-        raise AIError("下载 cmhub 图片失败: HTTP %s" % status)
-    chunks = []
-    total = 0
-    iterator = response.iter_content(chunk_size=65536) if hasattr(response, "iter_content") else [response.content]
-    for chunk in iterator:
-        if not chunk:
-            continue
-        total += len(chunk)
-        if total > max_bytes:
-            raise AIError("下载 cmhub 图片失败: 图片超过大小上限")
-        chunks.append(chunk)
-    return b"".join(chunks)
+    finally:
+        if response is not None and hasattr(response, "close"):
+            response.close()
 
 
 def _should_use_curl_for_cmhub_download(mode):
@@ -1846,6 +1878,7 @@ def _download_cmhub_image_with_curl(
     read_timeout,
     max_bytes,
     use_system_proxy=False,
+    should_stop=None,
 ):
     curl_path = _find_system_curl()
     if not curl_path:
@@ -1881,20 +1914,36 @@ def _download_cmhub_image_with_curl(
         ]
         if not bool(use_system_proxy):
             args.extend(["--noproxy", "*"])
+        process = None
         try:
-            completed = subprocess.run(
+            process = subprocess.Popen(
                 args,
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
-                timeout=max(2, int(connect_timeout) + int(read_timeout) + 10),
-                check=False,
                 shell=False,
                 **_subprocess_hidden_window_kwargs(),
             )
-        except (OSError, subprocess.TimeoutExpired) as exc:
+            deadline = time.monotonic() + max(
+                2,
+                int(connect_timeout) + int(read_timeout) + 10,
+            )
+            while process.poll() is None:
+                try:
+                    _raise_if_download_cancelled(should_stop)
+                except CancelledError:
+                    _stop_download_process(process)
+                    raise
+                if time.monotonic() >= deadline:
+                    _stop_download_process(process)
+                    raise AIError("下载 cmhub 图片失败: curl 执行超时")
+                time.sleep(0.1)
+            process.communicate()
+        except CancelledError:
+            raise
+        except OSError as exc:
             raise AIError("下载 cmhub 图片失败: curl 执行失败") from exc
-        if completed.returncode != 0:
-            raise AIError("下载 cmhub 图片失败: curl 退出码 %s" % completed.returncode)
+        if process.returncode != 0:
+            raise AIError("下载 cmhub 图片失败: curl 退出码 %s" % process.returncode)
         size = os.path.getsize(temp_output_path)
         if size > max_bytes:
             raise AIError("下载 cmhub 图片失败: 图片超过大小上限")
@@ -1909,6 +1958,37 @@ def _download_cmhub_image_with_curl(
                     pass
 
 
+def _stop_download_process(process):
+    if process is None or process.poll() is not None:
+        return
+    try:
+        process.terminate()
+        process.wait(timeout=2)
+    except (OSError, subprocess.TimeoutExpired):
+        try:
+            process.kill()
+            process.wait(timeout=2)
+        except (OSError, subprocess.TimeoutExpired):
+            pass
+
+
+def _raise_if_download_cancelled(should_stop):
+    try:
+        stopped = bool(should_stop and should_stop())
+    except Exception:
+        stopped = False
+    if stopped:
+        raise CancelledError()
+
+
+def _sleep_download_retry(delay_seconds, should_stop):
+    deadline = time.monotonic() + max(0.0, float(delay_seconds or 0.0))
+    while time.monotonic() < deadline:
+        _raise_if_download_cancelled(should_stop)
+        time.sleep(min(0.1, max(0.0, deadline - time.monotonic())))
+    _raise_if_download_cancelled(should_stop)
+
+
 def _curl_config_quote(value):
     text = str(value or "")
     return '"' + text.replace("\\", "\\\\").replace('"', '\\"') + '"'