cmshoppe/app/image_studio_generation.py

"""Multi-source generation orchestration for AI image studio jobs."""

from __future__ import annotations

import os
import threading
import time
import urllib.parse
import uuid
from concurrent.futures import CancelledError, FIRST_COMPLETED, ThreadPoolExecutor, wait

from . import ai, appconfig, image_studio
from .version import APP_VERSION


MAX_IMAGE_STUDIO_WORKERS = 5
_GLOBAL_IMAGE_STUDIO_SLOTS = threading.BoundedSemaphore(MAX_IMAGE_STUDIO_WORKERS)
CMHUB_IMAGE_STUDIO_MAX_INPUT_IMAGES = 8
CMHUB_IMAGE_STUDIO_MAX_SINGLE_INPUT_BYTES = 10 * 1024 * 1024
CMHUB_IMAGE_STUDIO_MAX_TOTAL_INPUT_BYTES = 32 * 1024 * 1024


class ImageStudioGenerationError(RuntimeError):
    """Raised when AI studio image generation cannot complete."""


class ImageStudioRuntimeLease:
    """A process-lifetime lock used only to make startup recovery safe."""

    def __init__(self, handle, path):
        self._handle = handle
        self.path = path
        self.session_id = "startup-" + uuid.uuid4().hex

    def release(self):
        handle, self._handle = self._handle, None
        if handle is None:
            return
        try:
            handle.seek(0)
            if os.name == "nt":
                import msvcrt

                msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1)
            else:  # pragma: no cover - Windows release is the supported path.
                import fcntl

                fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
        finally:
            handle.close()


def acquire_startup_recovery_lease(data_dir):
    """Acquire a non-blocking process lease; return None when another app owns it."""

    directory = os.path.abspath(str(data_dir or ""))
    if not directory:
        return None
    try:
        os.makedirs(directory, exist_ok=True)
        path = os.path.join(directory, ".image_studio_generation.lock")
        handle = open(path, "a+b")
        if not handle.read(1):
            handle.write(b"0")
            handle.flush()
        handle.seek(0)
        if os.name == "nt":
            import msvcrt

            msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1)
        else:  # pragma: no cover - retained for developer tests outside Windows.
            import fcntl

            fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
        return ImageStudioRuntimeLease(handle, path)
    except (OSError, ImportError):
        try:
            handle.close()
        except Exception:
            pass
        return None


def recover_stale_direct_jobs_at_startup(*, lease, path=None):
    """Recover interrupted direct jobs only while the process lease is held."""

    if lease is None:
        return {"recovered": [], "skipped": True}
    recovered = image_studio.fail_stale_direct_jobs(path=path)
    return {"recovered": recovered, "skipped": False}


RATIO_OUTPUT_SIZES = {
    "1:1": ("1024x1024", False),
    "3:4": ("1024x1536", True),
    "9:16": ("1024x1536", True),
    "4:3": ("1536x1024", True),
    "16:9": ("1536x1024", True),
}


def _notify(callback, event):
    if callback is None:
        return
    try:
        callback(dict(event))
    except Exception:
        pass


def _runtime(config, cmhub_config_path):
    return ai._cmhub_runtime(config, "image", cmhub_config_path)


def _direct_runtime(config):
    ai_cfg = appconfig.ai_config(config)
    models_path = appconfig.ai_models_config_path(config)
    model = ai._role_model(
        "image",
        ai_cfg.get("default_image_model"),
        models_path,
        models=ai._runtime_direct_models(config),
    )
    error = appconfig.image_model_config_error(model)
    if error:
        raise ImageStudioGenerationError(error)
    return {
        "model": model,
        "connect_timeout": ai._connect_timeout(model),
        "read_timeout": ai._read_timeout(model, config),
        "quality": ai._jpg_quality(ai_cfg.get("jpg_quality", 90)),
    }


def _is_default_gateway_job(job):
    return (
        str(getattr(job, "generation_source", "") or "").strip().lower()
        == image_studio.GENERATION_SOURCE_CMHUB
        and str(getattr(job, "provider", "") or "").strip().lower()
        == image_studio.PROVIDER_CMHUB
    )


def _is_direct_gateway_job(job):
    return (
        str(getattr(job, "generation_source", "") or "").strip().lower()
        == image_studio.GENERATION_SOURCE_DIRECT
        and str(getattr(job, "provider", "") or "").strip().lower()
        == image_studio.PROVIDER_OPENAI_IMAGES_EDITS
    )


def generation_source_for_config(config):
    """Return the persisted source/provider pair for one frozen run config."""

    if appconfig.ai_backend(config) == "direct":
        return {
            "generation_source": image_studio.GENERATION_SOURCE_DIRECT,
            "provider": image_studio.PROVIDER_OPENAI_IMAGES_EDITS,
        }
    return {
        "generation_source": image_studio.GENERATION_SOURCE_CMHUB,
        "provider": image_studio.PROVIDER_CMHUB,
    }


def requested_output_spec(aspect_ratio):
    """Map a product-suite ratio to an explicit request size and approximation flag."""

    ratio = str(aspect_ratio or "1:1")
    size, approximate = RATIO_OUTPUT_SIZES.get(ratio, RATIO_OUTPUT_SIZES["1:1"])
    return {
        "aspect_ratio": ratio,
        "requested_output_size": size,
        "approximate_ratio": bool(approximate),
    }


def _ensure_new_submission_allowed(config):
    if appconfig.ai_backend(config) not in {"cmhub", "direct"}:
        raise ImageStudioGenerationError("商品套图生成网关配置无效，请到⑤设置检查")


def _generated_kind(job_type):
    return "generated_detail" if str(job_type) == "detail" else "generated_main"


def _output_path(project, job, image_root):
    dirs = image_studio.project_image_dirs(image_root, project)
    filename = "job_%06d_%s.jpg" % (int(job.id), _generated_kind(job.job_type))
    return os.path.join(dirs["generated"], filename)


def _source_path(source_asset):
    path = str(getattr(source_asset, "local_path", "") or "").strip()
    if not path or not os.path.isfile(path):
        raise ImageStudioGenerationError("源图尚未下载到本地，不能生成")
    return path


def _build_cmhub_images(source_asset, reference_assets=()):
    """Build the ordered cmhub image input array from local image assets."""
    candidates = [source_asset] + list(reference_assets or [])
    selected = candidates[:CMHUB_IMAGE_STUDIO_MAX_INPUT_IMAGES]
    omitted_count = max(0, len(candidates) - len(selected))
    images = []
    total_bytes = 0
    for index, asset in enumerate(selected, 1):
        source_path = _source_path(asset)
        file_size = os.path.getsize(source_path)
        if file_size > CMHUB_IMAGE_STUDIO_MAX_SINGLE_INPUT_BYTES:
            role = "主图" if index == 1 else "参考图"
            raise ImageStudioGenerationError("商品套图%s超过10MiB，不能提交" % role)
        data_url = ai._image_data_url(source_path)
        total_bytes += len(data_url.encode("utf-8"))
        if total_bytes > CMHUB_IMAGE_STUDIO_MAX_TOTAL_INPUT_BYTES:
            raise ImageStudioGenerationError(
                "商品套图提交图片总大小超过32MiB，请减少图片或更换较小原图"
            )
        images.append({"image_base64": data_url})
    if not images:
        raise ImageStudioGenerationError("商品套图至少需要一张本地原图")
    return images, omitted_count


def create_generation_jobs(
    project_id,
    source_asset_id,
    prompt,
    count,
    *,
    job_type="main",
    config=None,
    path=None,
):
    _ensure_new_submission_allowed(config)
    total = max(0, int(count or 0))
    if total <= 0:
        raise ImageStudioGenerationError("生成数量必须大于0")
    source = generation_source_for_config(config)
    jobs = []
    for _ in range(total):
        jobs.append(
            image_studio.create_job(
                project_id,
                source_asset_id=source_asset_id,
                job_type=job_type,
                prompt=prompt,
                generation_source=source["generation_source"],
                provider=source["provider"],
                path=path,
            )
        )
    return jobs


def generate_image_jobs(
    project_id,
    source_asset_id,
    prompt,
    count,
    *,
    job_type="main",
    aspect_ratio="1:1",
    config=None,
    cmhub_config_path=None,
    path=None,
    should_stop=None,
    on_event=None,
    run_session_id=None,
):
    jobs = create_generation_jobs(
        project_id,
        source_asset_id,
        prompt,
        count,
        job_type=job_type,
        config=config,
        path=path,
    )
    return run_jobs(
        jobs,
        aspect_ratio=aspect_ratio,
        config=config,
        cmhub_config_path=cmhub_config_path,
        path=path,
        should_stop=should_stop,
        on_event=on_event,
        run_session_id=run_session_id,
    )


def resume_image_jobs(
    *,
    project_id=None,
    aspect_ratio="1:1",
    config=None,
    cmhub_config_path=None,
    path=None,
    should_stop=None,
    on_event=None,
    run_session_id=None,
):
    jobs = image_studio.list_resumable_jobs(
        path=path,
        project_id=project_id,
        include_failed_downloads=True,
    )
    return run_jobs(
        jobs,
        aspect_ratio=aspect_ratio,
        config=config,
        cmhub_config_path=cmhub_config_path,
        path=path,
        should_stop=should_stop,
        on_event=on_event,
        run_session_id=run_session_id,
    )


def run_jobs(
    jobs,
    *,
    aspect_ratio="1:1",
    config=None,
    cmhub_config_path=None,
    path=None,
    should_stop=None,
    on_event=None,
    run_session_id=None,
):
    cfg = appconfig.load_config() if config is None else config
    job_list = list(jobs or [])
    if not job_list:
        return {"total": 0, "success": 0, "failed": 0, "cancelled": 0, "jobs": []}
    rejected = []
    for job in job_list:
        if _is_direct_gateway_job(job) and getattr(job, "task_id", None):
            rejected.append(job)
        elif not _is_default_gateway_job(job) and not _is_direct_gateway_job(job):
            rejected.append(job)
    job_list = [job for job in job_list if job not in rejected]
    summary = {
        "total": len(job_list) + len(rejected),
        "success": 0,
        "failed": len(rejected),
        "cancelled": 0,
        "jobs": [
            {
                "job": job,
                "status": "failed",
                "error": "AI工场任务来源无效，不能执行",
            }
            for job in rejected
        ],
    }
    if not job_list:
        return summary

    ai_cfg = appconfig.ai_config(cfg)
    image_root = appconfig.image_dir(cfg)
    runtimes = {}
    if any(_is_default_gateway_job(job) for job in job_list):
        runtimes[image_studio.GENERATION_SOURCE_CMHUB] = _runtime(
            cfg,
            cmhub_config_path,
        )
    if any(_is_direct_gateway_job(job) for job in job_list):
        runtimes[image_studio.GENERATION_SOURCE_DIRECT] = _direct_runtime(cfg)
        run_session_id = str(run_session_id or "").strip() or "direct-" + uuid.uuid4().hex
    output_spec = requested_output_spec(aspect_ratio)
    max_workers = min(
        MAX_IMAGE_STUDIO_WORKERS,
        max(1, int(ai_cfg.get("image_concurrency", 1) or 1)),
        len(job_list),
    )
    should_stop = should_stop or (lambda: False)
    lock = threading.Lock()

    def record(result):
        with lock:
            summary["jobs"].append(result)
            if result.get("status") == "succeeded":
                summary["success"] += 1
            elif result.get("status") == "cancelled":
                summary["cancelled"] += 1
            else:
                summary["failed"] += 1

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(
                _run_one_job_with_global_slot,
                job.id,
                runtimes,
                cfg,
                image_root,
                output_spec,
                path,
                should_stop,
                on_event,
                run_session_id,
            ): job
            for job in job_list
        }
        while futures:
            done, _ = wait(
                set(futures),
                timeout=0.2,
                return_when=FIRST_COMPLETED,
            )
            for future in done:
                futures.pop(future)
                try:
                    record(future.result())
                except Exception as exc:
                    record({"status": "failed", "error": str(exc)})
            if should_stop():
                for future, job in list(futures.items()):
                    if future.cancel():
                        image_studio.update_job_status(
                            job.id,
                            "cancelled",
                            error="用户停止",
                            recovery_action=_recovery_action_for_job(job),
                            path=path,
                        )
                        futures.pop(future, None)
                        record({"job": job, "status": "cancelled", "error": "用户停止"})
    summary["output"] = dict(output_spec)
    return summary


def _run_one_job_with_global_slot(*args):
    with _GLOBAL_IMAGE_STUDIO_SLOTS:
        return _run_one_job(*args)


def _run_one_job(
    job_id,
    runtimes,
    config,
    image_root,
    output_spec,
    db_path,
    should_stop,
    on_event,
    run_session_id,
):
    job = image_studio.get_job(job_id, path=db_path)
    if job is None:
        raise ImageStudioGenerationError("AI工场生图任务不存在")
    if should_stop():
        updated = image_studio.update_job_status(
            job.id,
            "cancelled",
            error="用户停止",
            recovery_action=_recovery_action_for_job(job),
            path=db_path,
        )
        return {"job": updated, "status": "cancelled", "error": "用户停止"}
    project = image_studio.get_project(job.project_id, path=db_path)
    source_asset = image_studio.get_asset(job.source_asset_id, path=db_path)
    if project is None or source_asset is None:
        updated = image_studio.update_job_status(
            job.id,
            "failed",
            error="项目或源图不存在",
            recovery_action=image_studio.JOB_RECOVERY_REGENERATE,
            path=db_path,
        )
        return {"job": updated, "status": "failed", "error": "项目或源图不存在"}
    is_direct = _is_direct_gateway_job(job)
    if not (_is_default_gateway_job(job) or is_direct):
        raise ImageStudioGenerationError("AI工场任务来源无效，不能执行")
    runtime = runtimes.get(job.generation_source)
    if runtime is None:
        raise ImageStudioGenerationError("当前运行缺少该任务来源的配置快照")
    try:
        image_studio.update_job_status(
            job.id,
            "running",
            run_session_id=run_session_id if is_direct else None,
            path=db_path,
        )
        reference_assets = ()
        if not job.task_id:
            reference_assets = _reference_assets_for_job(job, db_path)
        out_path = _output_path(project, job, image_root)
        if is_direct:
            request_result = _submit_direct_job(
                job,
                source_asset,
                runtime,
                config,
                output_spec,
                should_stop,
                on_event,
                reference_assets=reference_assets,
            )
            # A submitted direct request cannot be cancelled. If bytes arrive after
            # stop was requested, save them before ending the remaining queue.
            saved_path = _save_direct_job_image(
                request_result,
                out_path,
                on_event,
                job.id,
            )
            remote_url = None
        else:
            request_result = _submit_or_resume_job(
                job,
                source_asset,
                runtime,
                config,
                output_spec["aspect_ratio"],
                db_path,
                should_stop,
                on_event,
                reference_assets=reference_assets,
            )
            image_studio.update_job_status(job.id, "running", path=db_path)
            request_result = _poll_job(
                job.id,
                request_result["task_id"],
                runtime,
                request_result,
                db_path,
                should_stop,
                on_event,
            )
            _raise_if_stopped(should_stop)
            request_result["resolution"] = output_spec["requested_output_size"]
            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:
                try:
                    if os.path.isfile(saved_path):
                        os.remove(saved_path)
                finally:
                    raise
            remote_url = request_result["image_url"]
        rendered_width, rendered_height = _rendered_image_size(saved_path)
        asset = image_studio.add_asset(
            project.id,
            _generated_kind(job.job_type),
            remote_url=remote_url,
            local_path=saved_path,
            aspect_ratio=output_spec["aspect_ratio"],
            requested_output_size=output_spec["requested_output_size"],
            rendered_width=rendered_width,
            rendered_height=rendered_height,
            parent_asset_id=source_asset.id,
            prompt=job.prompt,
            path=db_path,
        )
        updated = image_studio.update_job_status(
            job.id,
            "succeeded",
            output_asset_id=asset.id,
            points_balance=request_result.get("points_balance"),
            path=db_path,
        )
        _notify(on_event, {"job_id": job.id, "step": "job_done", "result": "success"})
        return {"job": updated, "asset": asset, "status": "succeeded"}
    except Exception as exc:
        cancelled = (
            not is_direct
            and (isinstance(exc, CancelledError) or "停止" in str(exc))
        )
        status = "cancelled" if cancelled else "failed"
        if cancelled:
            error = "用户已停止，已提交任务可稍后继续查询"
        elif is_direct:
            error = _direct_failure_message(exc)
        else:
            error = str(exc)
        current_job = image_studio.get_job(job.id, path=db_path)
        updated = image_studio.update_job_status(
            job.id,
            status,
            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": error,
            },
        )
        return {"job": updated, "status": status, "error": error}


def _reference_assets_for_job(job, db_path):
    assets = []
    for asset_id in image_studio.job_reference_asset_ids(job):
        asset = image_studio.get_asset(asset_id, path=db_path)
        if asset is None:
            raise ImageStudioGenerationError("参考图资产不存在，无法继续生成")
        try:
            _source_path(asset)
        except ImageStudioGenerationError as exc:
            raise ImageStudioGenerationError("参考图尚未下载到本地，无法继续生成") from exc
        assets.append(asset)
    return assets


def _direct_image_paths(source_asset, reference_assets=()):
    """Return ordered local paths; the first source image is the subject image."""

    return [_source_path(source_asset)] + [
        _source_path(asset) for asset in (reference_assets or ())
    ]


def _submit_direct_job(
    job,
    source_asset,
    runtime,
    config,
    output_spec,
    should_stop,
    on_event,
    reference_assets=(),
):
    if job.task_id:
        raise ImageStudioGenerationError("自定义网关任务不能继续查询，请手动重新生成")
    _raise_if_stopped(should_stop)
    image_paths = _direct_image_paths(source_asset, reference_assets)
    body, content_type = ai._image_edit_body(
        runtime["model"],
        job.prompt,
        image_paths,
        output_spec["requested_output_size"],
    )
    _notify(on_event, {"job_id": job.id, "step": "cover_submit", "result": "start"})
    # Direct images/edits has no idempotency or query endpoint. Retrying an
    # uncertain request could create and charge a second image, so one call only.
    data = ai._call_with_retry(
        runtime["model"],
        body,
        config,
        1,
        request_kind="multipart",
        content_type=content_type,
    )
    image_bytes = ai._extract_image_bytes(data, runtime["model"], config)
    _notify(on_event, {"job_id": job.id, "step": "cover_request", "result": "success"})
    return {
        "image_bytes": image_bytes,
        "resolution": output_spec["requested_output_size"],
        "quality": runtime["quality"],
    }


def _save_direct_job_image(request_result, out_path, on_event, job_id):
    _notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "start"})
    saved_path = ai._save_jpeg(
        request_result["image_bytes"],
        out_path,
        request_result["resolution"],
        request_result["quality"],
    )
    _notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "success"})
    return saved_path


def _rendered_image_size(path):
    try:
        from PIL import Image

        with Image.open(path) as image:
            width, height = image.size
        return int(width), int(height)
    except Exception as exc:
        raise ImageStudioGenerationError("生成图片保存后无法读取尺寸") from exc


def _direct_failure_message(exc):
    text = str(exc or "").lower()
    if "timeout" in text or "超时" in text:
        return "自定义网关生成超时，无法确认服务商是否已处理，请手动重新生成"
    return "自定义网关生成失败，结果无法确认，请检查图片模型配置后手动重新生成"


def _submit_or_resume_job(
    job,
    source_asset,
    runtime,
    config,
    aspect_ratio,
    db_path,
    should_stop,
    on_event,
    reference_assets=(),
):
    if job.task_id:
        if not _is_default_gateway_job(job):
            raise ImageStudioGenerationError("该已提交任务不属于默认网关，不能继续查询")
        _notify(on_event, {"job_id": job.id, "step": "cover_request", "result": "resume", "task_id": job.task_id})
        return _request_result(job.task_id, runtime, config)
    _ensure_new_submission_allowed(config)
    _raise_if_stopped(should_stop)
    ai_cfg = appconfig.ai_config(config)
    resolution = str(ai_cfg.get("resolution", "1k") or "1k")
    images, omitted_count = _build_cmhub_images(source_asset, reference_assets)
    if omitted_count:
        _notify(
            on_event,
            {
                "job_id": job.id,
                "step": "cover_submit",
                "result": "warning",
                "detail": "图片最多提交8张，已忽略%d张参考图" % omitted_count,
            },
        )
    payload = {
        "prompt": str(job.prompt or ""),
        "model": runtime["alias"],
        "images": images,
        "resolution": ai._normalize_cmhub_resolution(resolution),
        "aspect_ratio": str(aspect_ratio or "1:1"),
    }
    _notify(on_event, {"job_id": job.id, "step": "cover_submit", "result": "start"})
    data = ai._cmhub_call_with_retry(
        "POST",
        appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/image/tasks"),
        runtime["api_key"],
        payload=payload,
        connect_timeout=runtime["connect_timeout"],
        read_timeout=ai.CMHUB_IMAGE_SUBMIT_READ_TIMEOUT_SECONDS,
        attempts=max(1, int(ai_cfg.get("retry", 0) or 0) + 1),
        headers_extra={
            "Idempotency-Key": job.task_key,
            "X-Client-Version": str(APP_VERSION),
        },
    )
    image_task_id = str(data.get("task_id") or "").strip()
    if not image_task_id:
        raise ImageStudioGenerationError("cmhub 生图任务提交返回格式错误")
    submitted = image_studio.set_job_submitted(
        job.id,
        image_task_id,
        call_id=data.get("call_id"),
        points_cost=data.get("points_cost"),
        points_balance=data.get("points_balance"),
        path=db_path,
    )
    _notify(
        on_event,
        {
            "job_id": job.id,
            "step": "cover_submit",
            "result": "success",
            "task_id": submitted.task_id,
            "call_id": submitted.call_id,
            "points_cost": submitted.points_cost,
            "points_balance": submitted.points_balance,
        },
    )
    return _request_result(image_task_id, runtime, config, points_balance=data.get("points_balance"))


def _request_result(task_id, runtime, config, points_balance=None):
    ai_cfg = appconfig.ai_config(config)
    resolution = str(ai_cfg.get("resolution", "1k") or "1k")
    return {
        "task_id": task_id,
        "connect_timeout": runtime["connect_timeout"],
        "read_timeout": ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS,
        "resolution": resolution,
        "quality": ai._jpg_quality(ai_cfg.get("jpg_quality", 90)),
        "use_system_proxy": runtime["use_system_proxy"],
        "download_with_curl": runtime["download_with_curl"],
        "points_balance": points_balance,
    }


def _poll_job(job_id, task_id, runtime, request_result, db_path, should_stop, on_event):
    poll_url = appconfig.cmhub_request_url(
        runtime["base_url"],
        "/api/v1/generate/image/tasks/%s" % urllib.parse.quote(str(task_id), safe=""),
    )
    started = time.perf_counter()
    deadline = started + max(1, int(ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS))
    poll_index = 0
    while True:
        _raise_if_stopped(should_stop)
        if time.perf_counter() >= deadline:
            raise ImageStudioGenerationError("等待 cmhub 生图任务完成超时，下次可继续查询")
        _notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": "start", "task_id": task_id})
        data = ai._cmhub_call_once(
            "GET",
            poll_url,
            runtime["api_key"],
            payload=None,
            connect_timeout=runtime["connect_timeout"],
            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})
            ai._sleep_cmhub_poll(poll_index, should_stop)
            poll_index += 1
            continue
        if status == "succeeded":
            image_url = ai._extract_cmhub_image_url(data, runtime["base_url"])
            if not image_url:
                raise ImageStudioGenerationError("cmhub 生图任务成功但没有图片地址")
            merged = dict(request_result)
            merged["image_url"] = image_url
            merged["points_balance"] = data.get("points_balance", merged.get("points_balance"))
            _notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": "success", "task_id": task_id})
            return merged
        if status in {"failed", "expired"}:
            error = data.get("error") if isinstance(data.get("error"), dict) else {}
            message = str(error.get("message") or error.get("code") or status)
            image_studio.update_job_status(
                job_id,
                status,
                error=message,
                recovery_action=image_studio.JOB_RECOVERY_REGENERATE,
                path=db_path,
            )
            raise ImageStudioGenerationError(message)
        raise ImageStudioGenerationError("cmhub 生图任务状态返回格式错误")


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"})
    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


def _raise_if_stopped(should_stop):
    try:
        stopped = bool(should_stop and should_stop())
    except Exception:
        stopped = False
    if stopped:
        raise ImageStudioGenerationError("用户已停止，已提交任务可稍后继续查询")


def _recovery_action_for_job(job):
    if (
        job is not None
        and getattr(job, "task_id", None)
        and getattr(job, "recovery_action", None) == image_studio.JOB_RECOVERY_RESUME
    ):
        return image_studio.JOB_RECOVERY_RESUME
    return image_studio.JOB_RECOVERY_REGENERATE