feat(product-suite): add direct job state machine

app/image_studio.py

@@ -33,6 +33,14 @@ JOB_RECOVERY_ACTIONS = {
     JOB_RECOVERY_RESUME,
     JOB_RECOVERY_REGENERATE,
 }
+GENERATION_SOURCE_CMHUB = "cmhub"
+GENERATION_SOURCE_DIRECT = "direct"
+PROVIDER_CMHUB = "cmhub"
+PROVIDER_OPENAI_IMAGES_EDITS = "openai_images_edits"
+GENERATION_SOURCE_PROVIDERS = {
+    GENERATION_SOURCE_CMHUB: PROVIDER_CMHUB,
+    GENERATION_SOURCE_DIRECT: PROVIDER_OPENAI_IMAGES_EDITS,
+}
 SELECTION_TYPES = {"main", "detail"}
 
 
@@ -65,6 +73,9 @@ class ImageStudioAsset:
     remote_url: Optional[str]
     local_path: Optional[str]
     aspect_ratio: Optional[str]
+    requested_output_size: Optional[str]
+    rendered_width: Optional[int]
+    rendered_height: Optional[int]
     parent_asset_id: Optional[int]
     prompt: Optional[str]
     status: str
@@ -95,6 +106,7 @@ class ImageStudioJob:
     call_id: Optional[str]
     generation_round_key: Optional[str]
     generation_slot_index: Optional[int]
+    run_session_id: Optional[str]
     created_at: str
     updated_at: str
     submitted_at: Optional[str]
@@ -705,6 +717,9 @@ def add_asset(
     remote_url=None,
     local_path=None,
     aspect_ratio=None,
+    requested_output_size=None,
+    rendered_width=None,
+    rendered_height=None,
     parent_asset_id=None,
     prompt=None,
     status=ASSET_STATUS_AVAILABLE,
@@ -720,8 +735,9 @@ def add_asset(
                 """
                 INSERT INTO image_studio_assets
                   (project_id, kind, remote_url, local_path, aspect_ratio,
+                   requested_output_size, rendered_width, rendered_height,
                    parent_asset_id, prompt, status, source_order, created_at, updated_at)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                 """,
                 (
                     int(project_id),
@@ -729,6 +745,9 @@ def add_asset(
                     remote_url,
                     abs_local_path,
                     aspect_ratio,
+                    str(requested_output_size).strip() if requested_output_size else None,
+                    int(rendered_width) if rendered_width is not None else None,
+                    int(rendered_height) if rendered_height is not None else None,
                     parent_asset_id,
                     prompt,
                     str(status or ASSET_STATUS_AVAILABLE),
@@ -1161,6 +1180,11 @@ def create_job(
     path=None,
     conn=None,
 ):
+    source = str(generation_source or GENERATION_SOURCE_CMHUB).strip().lower()
+    expected_provider = GENERATION_SOURCE_PROVIDERS.get(source)
+    provider = str(provider or expected_provider or "").strip().lower()
+    if expected_provider is None or provider != expected_provider:
+        raise db.DbError("AI工场任务来源与服务商组合无效")
     now = _now()
     task_key = str(task_key or _task_key(project_id))
     reference_ids = _parse_reference_asset_ids(reference_asset_ids, source_asset_id)
@@ -1200,8 +1224,8 @@ def create_job(
                         int(project_id),
                         source_asset_id,
                         reference_json,
-                        str(generation_source or "cmhub"),
-                        str(provider or "cmhub"),
+                        source,
+                        provider,
                         str(job_type),
                         task_key,
                         str(prompt or ""),
@@ -1664,6 +1688,14 @@ def set_job_submitted(job_id, task_id, *, call_id=None, points_cost=None, points
     now = _now()
     with _connection(conn, path) as database:
         with database:
+            job = get_job(job_id, conn=database)
+            if job is None:
+                raise db.DbError("AI工场任务不存在")
+            if (
+                job.generation_source != GENERATION_SOURCE_CMHUB
+                or job.provider != PROVIDER_CMHUB
+            ):
+                raise db.DbError("只有默认网关任务可以记录查询任务编号")
             database.execute(
                 """
                 UPDATE image_studio_jobs
@@ -1700,6 +1732,7 @@ def update_job_status(
     points_balance=None,
     recovery_action=None,
     increment_attempts=False,
+    run_session_id=None,
     path=None,
     conn=None,
 ):
@@ -1726,6 +1759,7 @@ def update_job_status(
                     recovery_action = COALESCE(?, recovery_action),
                     output_asset_id = COALESCE(?, output_asset_id),
                     points_balance = COALESCE(?, points_balance),
+                    run_session_id = COALESCE(?, run_session_id),
                     attempts = attempts + ?,
                     finished_at = CASE WHEN ? THEN ? ELSE finished_at END,
                     updated_at = ?
@@ -1737,6 +1771,7 @@ def update_job_status(
                     recovery_action,
                     output_asset_id,
                     points_balance,
+                    str(run_session_id).strip() if run_session_id else None,
                     1 if increment_attempts else 0,
                     1 if terminal else 0,
                     now,
@@ -1751,13 +1786,35 @@ def list_resumable_jobs(path=None, conn=None, project_id=None, include_failed_do
     if include_failed_downloads:
         clauses = [
             "task_id IS NOT NULL",
+            "generation_source = ?",
+            "provider = ?",
             "recovery_action = ?",
             "status IN (?, ?, ?, ?)",
         ]
-        params = [JOB_RECOVERY_RESUME, "submitted", "running", "failed", "cancelled"]
+        params = [
+            GENERATION_SOURCE_CMHUB,
+            PROVIDER_CMHUB,
+            JOB_RECOVERY_RESUME,
+            "submitted",
+            "running",
+            "failed",
+            "cancelled",
+        ]
     else:
-        clauses = ["status IN (?, ?)", "task_id IS NOT NULL", "recovery_action = ?"]
-        params = ["submitted", "running", JOB_RECOVERY_RESUME]
+        clauses = [
+            "status IN (?, ?)",
+            "task_id IS NOT NULL",
+            "generation_source = ?",
+            "provider = ?",
+            "recovery_action = ?",
+        ]
+        params = [
+            "submitted",
+            "running",
+            GENERATION_SOURCE_CMHUB,
+            PROVIDER_CMHUB,
+            JOB_RECOVERY_RESUME,
+        ]
     if project_id is not None:
         clauses.append("project_id = ?")
         params.append(int(project_id))
@@ -1767,6 +1824,58 @@ def list_resumable_jobs(path=None, conn=None, project_id=None, include_failed_do
         return _fetch_all(database, sql, params, ImageStudioJob)
 
 
+def fail_stale_direct_jobs(*, active_run_session_ids=(), path=None, conn=None):
+    """Fail only confirmed stale synchronous direct jobs during application startup."""
+
+    active_ids = {
+        str(value).strip()
+        for value in (active_run_session_ids or ())
+        if str(value).strip()
+    }
+    with _connection(conn, path) as database:
+        rows = database.execute(
+            """
+            SELECT * FROM image_studio_jobs
+            WHERE generation_source = ?
+              AND provider = ?
+              AND status = 'running'
+            ORDER BY id
+            """,
+            (GENERATION_SOURCE_DIRECT, PROVIDER_OPENAI_IMAGES_EDITS),
+        ).fetchall()
+        stale_ids = [
+            int(row["id"])
+            for row in rows
+            if str(row["run_session_id"] or "").strip() not in active_ids
+        ]
+        if not stale_ids:
+            return []
+        now = _now()
+        with database:
+            database.executemany(
+                """
+                UPDATE image_studio_jobs
+                SET status = 'failed',
+                    error = ?,
+                    recovery_action = ?,
+                    finished_at = ?,
+                    updated_at = ?
+                WHERE id = ?
+                """,
+                [
+                    (
+                        "程序中断，无法确认生成结果，请手动重新生成",
+                        JOB_RECOVERY_REGENERATE,
+                        now,
+                        now,
+                        job_id,
+                    )
+                    for job_id in stale_ids
+                ],
+            )
+        return [get_job(job_id, conn=database) for job_id in stale_ids]
+
+
 def replace_selections(project_id, selection_type, asset_ids: Iterable[int], path=None, conn=None):
     selection = str(selection_type)
     if selection not in SELECTION_TYPES: