From dcaf68fd0779ad6c6a6996248265546daab151c1 Mon Sep 17 00:00:00 2001 From: chengma Date: Sat, 11 Jul 2026 14:29:58 +0800 Subject: [PATCH] feat(ai-studio): resume hosted image jobs --- app/gui/__init__.py | 1 + app/gui/tabs/image_studio.py | 57 ++++++++++++++++++++++- app/gui/workers.py | 47 +++++++++++++++++++ app/image_studio.py | 13 ++++-- app/image_studio_generation.py | 6 ++- docs/routes.md | 1 + docs/tasks/T-594.md | 10 ++++- tests/test_gui.py | 47 +++++++++++++++++++ tests/test_image_studio_generation.py | 65 +++++++++++++++++++++++++++ 9 files changed, 240 insertions(+), 7 deletions(-) diff --git a/app/gui/__init__.py b/app/gui/__init__.py index 3821b40..9c4d9e2 100644 --- a/app/gui/__init__.py +++ b/app/gui/__init__.py @@ -28,6 +28,7 @@ if QT_IMPORT_ERROR is None: ImageStudioExportWorker, ImageStudioGenerateJobsWorker, ImageStudioPullImagesWorker, + ImageStudioResumeJobsWorker, WriteBackWorker, ) from .tabs.accounts import AccountDialog, AccountsTab diff --git a/app/gui/tabs/image_studio.py b/app/gui/tabs/image_studio.py index 86022ec..936296f 100644 --- a/app/gui/tabs/image_studio.py +++ b/app/gui/tabs/image_studio.py @@ -16,6 +16,7 @@ from ..workers import ( from ..workers import ImageStudioExportWorker as _RealImageStudioExportWorker from ..workers import ImageStudioGenerateJobsWorker as _RealImageStudioGenerateJobsWorker from ..workers import ImageStudioPullImagesWorker as _RealImageStudioPullImagesWorker +from ..workers import ImageStudioResumeJobsWorker as _RealImageStudioResumeJobsWorker ASSET_MIME_TYPE = "application/x-cmshopee-image-studio-asset" @@ -57,6 +58,15 @@ def ImageStudioExportWorker(*args, **kwargs): ) +def ImageStudioResumeJobsWorker(*args, **kwargs): + return _call_package_attr( + "ImageStudioResumeJobsWorker", + _RealImageStudioResumeJobsWorker, + *args, + **kwargs, + ) + + def _drop_event_position(event): if hasattr(event, "position"): return event.position().toPoint() @@ -432,10 +442,13 @@ class ImageStudioTab(QWidget): action_layout = QHBoxLayout() self.start_button = QPushButton("开始生成") self.start_button.setObjectName("imageStudioStartButton") + self.resume_button = QPushButton("继续查询任务") + self.resume_button.setObjectName("imageStudioResumeButton") self.stop_button = QPushButton("停止") self.stop_button.setObjectName("imageStudioStopButton") self.stop_button.setEnabled(False) action_layout.addWidget(self.start_button) + action_layout.addWidget(self.resume_button) action_layout.addWidget(self.stop_button) layout.addLayout(action_layout) @@ -512,6 +525,7 @@ class ImageStudioTab(QWidget): self.template_delete_button.clicked.connect(self.delete_template) self.prompt_edit.textChanged.connect(self._save_project_prompt) self.start_button.clicked.connect(self.start_generation) + self.resume_button.clicked.connect(self.resume_generation_jobs) self.stop_button.clicked.connect(self.stop_generation) self.export_button.clicked.connect(self.export_selections) @@ -730,7 +744,7 @@ class ImageStudioTab(QWidget): values = [ "任务", "-", - self.JOB_STATUS_LABELS.get(obj.status, obj.status), + _job_status_text(obj, self.JOB_STATUS_LABELS), f"源图 #{obj.source_asset_id or '-'}", obj.error or "", ] @@ -1055,6 +1069,35 @@ class ImageStudioTab(QWidget): self._append_log("[AI工场] 已请求停止，正在等待安全边界") self._status("AI工场生成已请求停止", "warning") + def resume_generation_jobs(self, checked=False): + if self.current_project is None: + self._message("项目未打开", "请先打开一个AI工场项目。") + return + resumable = image_studio.list_resumable_jobs( + path=self.db_path, + project_id=self.current_project.id, + include_failed_downloads=True, + ) + if not resumable: + self._status("当前项目没有可继续查询的 cmhub 生图任务", "muted") + self._message("没有可继续查询任务", "当前项目没有已提交、生成中或下载失败的 cmhub 生图任务。") + return + self.progress_bar.setRange(0, len(resumable)) + self.progress_bar.setValue(0) + self._append_log(f"[AI工场] 继续查询 {len(resumable)} 个已提交任务") + worker = ImageStudioResumeJobsWorker( + project_id=self.current_project.id, + aspect_ratio=self.aspect_combo.currentData(), + db_path=self.db_path, + config=self.config, + cmhub_config_path=self.cmhub_config_path, + ) + worker.progress.connect(self._on_generate_progress) + worker.log.connect(self._append_log) + worker.finished.connect(self._on_generation_finished) + worker.failed.connect(self._on_worker_failed) + self._start_worker(worker, "AI工场继续查询") + def export_selections(self, checked=False): if self.current_project is None: self._message("项目未打开", "请先打开一个AI工场项目。") @@ -1203,6 +1246,7 @@ class ImageStudioTab(QWidget): self.count_spin.setEnabled(not running) self.aspect_combo.setEnabled(not running) self.start_button.setEnabled(not running) + self.resume_button.setEnabled(not running) self.stop_button.setEnabled(running) self.export_button.setEnabled(not running) @@ -1380,3 +1424,14 @@ def _selection_tooltip(selection_type, asset): f"{_selection_label(selection_type)}：{_asset_badge(getattr(asset, 'kind', ''))} " f"#{getattr(asset, 'id', '')}，比例 {getattr(asset, 'aspect_ratio', None) or '未知'}" ) + + +def _job_status_text(job, labels): + parts = [labels.get(job.status, job.status)] + if job.points_cost is not None: + parts.append(f"扣点{job.points_cost}") + if job.points_balance is not None: + parts.append(f"余额{job.points_balance}") + if job.call_id: + parts.append(f"call_id={job.call_id}") + return "，".join(parts) diff --git a/app/gui/workers.py b/app/gui/workers.py index c4650f2..42dba76 100644 --- a/app/gui/workers.py +++ b/app/gui/workers.py @@ -192,6 +192,53 @@ class ImageStudioGenerateJobsWorker(BaseWorker): return summary +class ImageStudioResumeJobsWorker(BaseWorker): + """Resume submitted/running or failed-download AI studio jobs.""" + + def __init__( + self, + *, + project_id=None, + aspect_ratio="1:1", + db_path=None, + config=None, + cmhub_config_path=None, + ): + super().__init__() + self.project_id = int(project_id) if project_id is not None else None + self.aspect_ratio = str(aspect_ratio or "1:1") + self.db_path = db_path + self.config = config + self.cmhub_config_path = cmhub_config_path + self._done = 0 + self._failed = 0 + self._lock = threading.Lock() + + def execute(self): + def on_event(payload): + event = dict(payload or {}) + self.log.emit(_format_image_studio_event(event)) + if event.get("step") == "job_done": + with self._lock: + self._done += 1 + if event.get("result") not in {"success"}: + self._failed += 1 + progress = {"done": self._done, "failed": self._failed} + self.progress.emit(progress) + + summary = image_studio_generation.resume_image_jobs( + project_id=self.project_id, + aspect_ratio=self.aspect_ratio, + config=self.config, + cmhub_config_path=self.cmhub_config_path, + path=self.db_path, + should_stop=self.should_cancel, + on_event=on_event, + ) + summary["project_id"] = self.project_id + return summary + + class ImageStudioExportWorker(BaseWorker): """Export AI studio final selections to local JPEG files.""" diff --git a/app/image_studio.py b/app/image_studio.py index ae7fb2c..d510a5d 100644 --- a/app/image_studio.py +++ b/app/image_studio.py @@ -738,9 +738,16 @@ def update_job_status( return get_job(job_id, conn=database) -def list_resumable_jobs(path=None, conn=None, project_id=None): - clauses = ["status IN (?, ?)", "task_id IS NOT NULL"] - params = ["submitted", "running"] +def list_resumable_jobs(path=None, conn=None, project_id=None, include_failed_downloads=False): + if include_failed_downloads: + clauses = [ + "task_id IS NOT NULL", + "(status IN (?, ?) OR (status = ? AND output_asset_id IS NULL))", + ] + params = ["submitted", "running", "failed"] + else: + clauses = ["status IN (?, ?)", "task_id IS NOT NULL"] + params = ["submitted", "running"] if project_id is not None: clauses.append("project_id = ?") params.append(int(project_id)) diff --git a/app/image_studio_generation.py b/app/image_studio_generation.py index a294005..0013136 100644 --- a/app/image_studio_generation.py +++ b/app/image_studio_generation.py @@ -120,7 +120,11 @@ def resume_image_jobs( should_stop=None, on_event=None, ): - jobs = image_studio.list_resumable_jobs(path=path, project_id=project_id) + jobs = image_studio.list_resumable_jobs( + path=path, + project_id=project_id, + include_failed_downloads=True, + ) return run_jobs( jobs, aspect_ratio=aspect_ratio, diff --git a/docs/routes.md b/docs/routes.md index 0662264..32a0acd 100644 --- a/docs/routes.md +++ b/docs/routes.md @@ -196,6 +196,7 @@ - 照片池展示原图、生成主图、生成详情图和在途/失败任务状态；单击可用图片设为源图，双击打开大图；右键移除只删除未被任务或终选引用的照片池记录，不删除本地图片文件。 - 右侧只有一个完整提示词框；模板目录固定为 `data/prompts/image_studio/`，与②标题/封面模板隔离。界面不显示“主提示词 / 每张动作词”。 - 生图固定走 cmhub 托管模型，使用⑤设置里的 cmhub Base URL/API Key/生图别名和图片并发；界面只显示扣点、余额、进度、失败，不展示自定义 Provider、API Key、生成来源选择或“导入本地图片”入口。 +- 「继续查询任务」会恢复当前项目已提交、生成中或下载失败但已有 `task_id` 的 cmhub 生图任务；恢复时只 poll/download 原任务，不再次 submit，不重复扣点。照片池中的任务行显示排队/生成/失败/过期/停止状态，并附带扣点、余额和 `call_id`，便于运营和技术排障。 - 底部终选盘分为主图和详情图两列；照片池中已下载/已生成且本地文件可用的图片可拖入终选，落到已有位置时按插入顺延，同一类别内同一照片只能出现一次，主图和详情图之间允许复用同一照片。 - 终选列表内可拖动重排，Delete 或右键「移出终选」只移出终选，不删除照片池资产或本地文件；拖放/移出失败时刷新回 SQLite 中的持久化顺序。 - 主图推荐 1:1；比例不匹配只用黄色轻提示和 tooltip 提醒，不硬拦。文件缺失或尚未下载的照片不能拖入终选。 diff --git a/docs/tasks/T-594.md b/docs/tasks/T-594.md index ce5f4d9..0212eff 100644 --- a/docs/tasks/T-594.md +++ b/docs/tasks/T-594.md @@ -3,7 +3,7 @@ id: T-594 title: AI工场 cmhub 托管生图故障恢复、计费一致性与运行回归 phase: 7 deps: [T-590, T-591] -status: TODO +status: DONE created: 2026-07-11 --- @@ -36,4 +36,10 @@ AI工场会同时持有多张 cmhub 异步生图任务。基础编排和 GUI 接 ## 执行记录 -（完成后记录故障矩阵、状态决策与验证结果。） +- 2026-07-11：完成 AI工场 cmhub 托管生图恢复与状态收口。 + - `image_studio.list_resumable_jobs(..., include_failed_downloads=True)` 支持把 failed 但已有 `task_id` 且无输出 asset 的任务纳入恢复队列，用于下载失败/本地保存失败后的继续查询。 + - `image_studio_generation.resume_image_jobs()` 改用恢复队列：已有 `task_id` 的任务只 poll/download，不再 POST submit，不重复创建 job 或重复扣点。 + - 新增 `ImageStudioResumeJobsWorker` 和 ⑥ AI工场「继续查询任务」按钮；无可续查任务时中文提示，有任务时复用现有进度/日志/刷新逻辑。 + - 照片池任务行显示 submitted/running/failed/expired/cancelled 中文状态，并补充扣点、余额、`call_id`，便于计费和失败排障。 + - 更新 `docs/routes.md`，明确继续查询不重复 submit、不重复扣点以及任务状态展示口径。 + - 验证：主工作区 targeted 测试通过；因无关默认提示词脏文件仍会影响全量 unittest，已在干净 worktree 仅套用 T-594 diff 后运行 `python -m ruff check app tests main.py`、`py -3.10 -m compileall app main.py`、`py -3.10 -m unittest discover -s tests`（369 tests）和 `git diff --check`，全部通过。 diff --git a/tests/test_gui.py b/tests/test_gui.py index 7b508cd..775619d 100644 --- a/tests/test_gui.py +++ b/tests/test_gui.py @@ -602,6 +602,53 @@ class GuiTests(TempDirMixin, unittest.TestCase): self.assert_removed(temp_dir) + def test_image_studio_tab_shows_resume_and_job_billing_status(self): + with self.make_temp_dir() as temp_dir: + cfg = self.make_config(temp_dir) + db.init_db(cfg["db_path"]) + project = image_studio.create_or_get_project( + account_alias="alias-a", + account_slug="alias_a", + item_id="51100639510", + path=cfg["db_path"], + ) + source = image_studio.add_asset( + project.id, + "original", + local_path=self.write_test_image(os.path.join(temp_dir, "source.jpg")), + path=cfg["db_path"], + ) + job = image_studio.create_job( + project.id, + source_asset_id=source.id, + path=cfg["db_path"], + ) + image_studio.set_job_submitted( + job.id, + "cmhub-task-1", + call_id="call-1", + points_cost=2, + points_balance=88, + path=cfg["db_path"], + ) + tab = ImageStudioTab(config=cfg, db_path=cfg["db_path"]) + self.addCleanup(tab.close) + tab._select_project(project.id) + + self.assertEqual("继续查询任务", tab.resume_button.text()) + statuses = [ + tab.pool_table.item(row, 2).text() + for row in range(tab.pool_table.rowCount()) + if tab.pool_table.item(row, 0).text() == "任务" + ] + self.assertEqual(1, len(statuses)) + self.assertIn("已提交", statuses[0]) + self.assertIn("扣点2", statuses[0]) + self.assertIn("余额88", statuses[0]) + self.assertIn("call_id=call-1", statuses[0]) + + self.assert_removed(temp_dir) + def test_startup_update_gate_forced_blocks_and_opens_download(self): boxes = [] diff --git a/tests/test_image_studio_generation.py b/tests/test_image_studio_generation.py index b55bfed..ed7d6f4 100644 --- a/tests/test_image_studio_generation.py +++ b/tests/test_image_studio_generation.py @@ -231,6 +231,71 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase): self.assert_removed(temp_dir) + def test_resume_failed_download_job_polls_without_new_submit(self): + with self.make_temp_dir() as temp_dir: + cfg, project, source = self._project_source(temp_dir) + job = image_studio.create_job( + project.id, + source_asset_id=source.id, + job_type="main", + prompt="恢复下载", + path=cfg["db_path"], + ) + image_studio.set_job_submitted( + job.id, + "cmhub-task-download", + call_id="call-download", + points_cost=2, + points_balance=88, + path=cfg["db_path"], + ) + image_studio.update_job_status( + job.id, + "failed", + error="下载新封面失败", + path=cfg["db_path"], + ) + + resumable = image_studio.list_resumable_jobs( + path=cfg["db_path"], + project_id=project.id, + include_failed_downloads=True, + ) + self.assertEqual([job.id], [item.id for item in resumable]) + + def fake_poll(method, url, api_key, **kwargs): + self.assertEqual("GET", method) + return { + "task_id": "cmhub-task-download", + "status": "succeeded", + "result": {"image_url": "https://cdn.example.com/recovered.png"}, + "points_balance": 88, + } + + with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \ + mock.patch("app.image_studio_generation.ai._cmhub_call_with_retry") as submit, \ + mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \ + mock.patch( + "app.image_studio_generation.ai._download_cmhub_image_with_retry", + return_value=(self._png_bytes(), 0.1), + ): + summary = image_studio_generation.resume_image_jobs( + project_id=project.id, + config=cfg, + path=cfg["db_path"], + ) + + self.assertEqual(1, summary["success"]) + submit.assert_not_called() + updated = image_studio.get_job(job.id, path=cfg["db_path"]) + self.assertEqual("succeeded", updated.status) + self.assertEqual("cmhub-task-download", updated.task_id) + self.assertEqual("call-download", updated.call_id) + assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"]) + self.assertEqual(1, len(assets)) + + self.assert_removed(temp_dir) + def test_failed_cmhub_task_marks_only_that_job_failed(self): with self.make_temp_dir() as temp_dir: cfg, project, source = self._project_source(temp_dir)