feat: add task workflow and asset downloads

2026-05-29 12:32:02 +08:00
parent f9c3393f84
commit 63e62d444c
61 changed files with 2773 additions and 2181 deletions
--- a/lib/server/data-store.ts
+++ b/lib/server/data-store.ts
@@ -1,7 +1,7 @@
 import { readFile, rename, writeFile } from "node:fs/promises";
 import { join } from "node:path";
 import { createClient, type SupabaseClient } from "@supabase/supabase-js";
-import type { AppState, Asset, GenerationJob, Project, UsageEvent } from "@/lib/types";
+import type { AppState, Asset, GenerationCapability, GenerationJob, GenerationStatus, Project, UsageEvent } from "@/lib/types";
 import { createId } from "@/lib/server/ids";
 import { dataDir, DEFAULT_OWNER_ID, ensureRuntimeDirs } from "@/lib/server/runtime";

@@ -12,6 +12,21 @@ type AssetInput = Omit<Asset, "id" | "createdAt" | "updatedAt"> & Partial<Pick<A
 type JobInput = Omit<GenerationJob, "id" | "createdAt" | "updatedAt"> & Partial<Pick<GenerationJob, "id" | "createdAt" | "updatedAt">>;
 type UsageInput = Omit<UsageEvent, "id" | "createdAt"> & Partial<Pick<UsageEvent, "id" | "createdAt">>;

+export type GenerationJobListFilters = {
+  ownerId?: string;
+  externalClientId?: string;
+  status?: GenerationStatus;
+  capability?: GenerationCapability;
+  limit?: number;
+  before?: string;
+};
+
+export type ClaimGenerationJobsInput = {
+  workerId: string;
+  limit?: number;
+  lockTimeoutMs?: number;
+};
+
 export async function listAssets(ownerId = DEFAULT_OWNER_ID): Promise<Asset[]> {
  const supabase = getSupabaseAdmin();
  if (supabase) {
@@ -82,19 +97,37 @@ export async function deleteAsset(id: string): Promise<Asset | null> {
 }

 export async function listGenerationJobs(ownerId = DEFAULT_OWNER_ID, limit = 200): Promise<GenerationJob[]> {
+  return listGenerationJobsFiltered({ ownerId, limit });
+}
+
+export async function listGenerationJobsFiltered(filters: GenerationJobListFilters = {}): Promise<GenerationJob[]> {
+  const ownerId = filters.ownerId || DEFAULT_OWNER_ID;
+  const limit = filters.limit || 200;
  const supabase = getSupabaseAdmin();
  if (supabase) {
-    const { data, error } = await supabase
+    let query = supabase
      .from("generation_jobs")
      .select("*")
      .eq("owner_id", ownerId)
      .order("created_at", { ascending: false })
      .limit(limit);
+    if (filters.externalClientId) query = query.eq("external_client_id", filters.externalClientId);
+    if (filters.status) query = query.eq("status", filters.status);
+    if (filters.capability) query = query.eq("capability", filters.capability);
+    if (filters.before) query = query.lt("created_at", filters.before);
+    const { data, error } = await query;
    if (error) throw new Error(error.message);
    return (data || []).map(jobFromRow);
  }
  const state = await readState();
-  return state.generationJobs.filter((job) => job.ownerId === ownerId).sort(sortNewest).slice(0, limit);
+  return state.generationJobs
+    .filter((job) => job.ownerId === ownerId)
+    .filter((job) => !filters.externalClientId || job.externalClientId === filters.externalClientId)
+    .filter((job) => !filters.status || job.status === filters.status)
+    .filter((job) => !filters.capability || job.capability === filters.capability)
+    .filter((job) => !filters.before || job.createdAt < filters.before)
+    .sort(sortNewest)
+    .slice(0, limit);
 }

 export async function getGenerationJob(id: string): Promise<GenerationJob | null> {
@@ -118,6 +151,11 @@ export async function createGenerationJob(input: JobInput): Promise<GenerationJo
    inputUrls: input.inputUrls || [],
    outputAssetIds: input.outputAssetIds || [],
    requestPayload: input.requestPayload || {},
+    priority: input.priority ?? 0,
+    attempts: input.attempts ?? 0,
+    maxAttempts: input.maxAttempts ?? 3,
+    scheduledAt: input.scheduledAt || now,
+    webhookAttempts: input.webhookAttempts ?? 0,
    createdAt: input.createdAt || now,
    updatedAt: input.updatedAt || now
  };
@@ -133,6 +171,100 @@ export async function createGenerationJob(input: JobInput): Promise<GenerationJo
  });
 }

+export async function findGenerationJobByIdempotency(
+  externalClientId: string,
+  idempotencyKey: string,
+  ownerId = DEFAULT_OWNER_ID
+): Promise<GenerationJob | null> {
+  const supabase = getSupabaseAdmin();
+  if (supabase) {
+    const { data, error } = await supabase
+      .from("generation_jobs")
+      .select("*")
+      .eq("owner_id", ownerId)
+      .eq("external_client_id", externalClientId)
+      .eq("idempotency_key", idempotencyKey)
+      .maybeSingle();
+    if (error) throw new Error(error.message);
+    return data ? jobFromRow(data) : null;
+  }
+  const state = await readState();
+  return state.generationJobs.find((job) => (
+    job.ownerId === ownerId &&
+    job.externalClientId === externalClientId &&
+    job.idempotencyKey === idempotencyKey
+  )) || null;
+}
+
+export async function claimGenerationJobs(input: ClaimGenerationJobsInput): Promise<GenerationJob[]> {
+  const limit = Math.max(1, Math.min(input.limit || 1, 20));
+  const lockTimeoutMs = input.lockTimeoutMs ?? 5 * 60 * 1000;
+  const supabase = getSupabaseAdmin();
+  if (supabase) {
+    const { data, error } = await supabase.rpc("claim_generation_jobs", {
+      p_worker_id: input.workerId,
+      p_limit: limit,
+      p_lock_timeout_seconds: Math.ceil(lockTimeoutMs / 1000)
+    });
+    if (error) throw new Error(`claim_generation_jobs failed: ${error.message}`);
+    return (Array.isArray(data) ? data : []).map(jobFromRow);
+  }
+
+  return mutateLocalState((state) => {
+    const now = new Date();
+    const nowIso = now.toISOString();
+    const staleBefore = new Date(now.getTime() - lockTimeoutMs).toISOString();
+    const selected = state.generationJobs
+      .filter((job) => isClaimableJob(job, nowIso, staleBefore))
+      .sort(sortClaimableJobs)
+      .slice(0, limit);
+    for (const job of selected) {
+      job.lockedAt = nowIso;
+      job.lockedBy = input.workerId;
+      if (!job.startedAt) job.startedAt = nowIso;
+      job.updatedAt = nowIso;
+    }
+    return selected.map((job) => ({ ...job }));
+  });
+}
+
+export async function clearGenerationJobLock(
+  id: string,
+  patch: Partial<GenerationJob> = {},
+  options: { clearProviderTaskId?: boolean } = {}
+): Promise<GenerationJob> {
+  const updatedAt = new Date().toISOString();
+  const supabase = getSupabaseAdmin();
+  if (supabase) {
+    const { data, error } = await supabase
+      .from("generation_jobs")
+      .update({
+        ...jobToRow({ ...patch, updatedAt } as GenerationJob),
+        locked_at: null,
+        locked_by: null,
+        ...(options.clearProviderTaskId ? { provider_task_id: null } : {})
+      })
+      .eq("id", id)
+      .select("*")
+      .single();
+    if (error) throw new Error(error.message);
+    return jobFromRow(data);
+  }
+  return mutateLocalState((state) => {
+    const index = state.generationJobs.findIndex((job) => job.id === id);
+    if (index === -1) throw new Error(`Generation job not found: ${id}`);
+    state.generationJobs[index] = {
+      ...state.generationJobs[index],
+      ...patch,
+      lockedAt: undefined,
+      lockedBy: undefined,
+      ...(options.clearProviderTaskId ? { providerTaskId: undefined } : {}),
+      updatedAt
+    };
+    return state.generationJobs[index];
+  });
+}
+
 export async function updateGenerationJob(id: string, patch: Partial<GenerationJob>): Promise<GenerationJob> {
  const updatedAt = new Date().toISOString();
  const supabase = getSupabaseAdmin();
@@ -252,6 +384,20 @@ function sortNewest<T extends { createdAt: string }>(a: T, b: T): number {
  return b.createdAt.localeCompare(a.createdAt);
 }

+function isClaimableJob(job: GenerationJob, nowIso: string, staleBefore: string): boolean {
+  if (["succeeded", "failed", "expired", "cancelled"].includes(job.status)) return false;
+  if ((job.scheduledAt || job.createdAt) > nowIso) return false;
+  return !job.lockedAt || job.lockedAt < staleBefore;
+}
+
+function sortClaimableJobs(a: GenerationJob, b: GenerationJob): number {
+  const priority = (b.priority || 0) - (a.priority || 0);
+  if (priority !== 0) return priority;
+  const scheduled = (a.scheduledAt || a.createdAt).localeCompare(b.scheduledAt || b.createdAt);
+  if (scheduled !== 0) return scheduled;
+  return a.createdAt.localeCompare(b.createdAt);
+}
+
 function assetToRow(asset: Partial<Asset>) {
  return {
    id: asset.id,
@@ -288,6 +434,7 @@ function jobToRow(job: Partial<GenerationJob>) {
  const row: Record<string, unknown> = {};
  if (job.id !== undefined) row.id = job.id;
  if (job.ownerId !== undefined) row.owner_id = job.ownerId;
+  if (job.externalClientId !== undefined) row.external_client_id = job.externalClientId;
  if (job.capability !== undefined) row.capability = job.capability;
  if (job.provider !== undefined) row.provider = job.provider;
  if (job.reqKey !== undefined) row.req_key = job.reqKey;
@@ -301,6 +448,19 @@ function jobToRow(job: Partial<GenerationJob>) {
  if (job.responsePayload !== undefined) row.response_payload = job.responsePayload;
  if (job.error !== undefined) row.error = job.error;
  if (job.retryOf !== undefined) row.retry_of = job.retryOf;
+  if (job.idempotencyKey !== undefined) row.idempotency_key = job.idempotencyKey;
+  if (job.idempotencyFingerprint !== undefined) row.idempotency_fingerprint = job.idempotencyFingerprint;
+  if (job.priority !== undefined) row.priority = job.priority;
+  if (job.attempts !== undefined) row.attempts = job.attempts;
+  if (job.maxAttempts !== undefined) row.max_attempts = job.maxAttempts;
+  if (job.scheduledAt !== undefined) row.scheduled_at = job.scheduledAt;
+  if (job.lockedAt !== undefined) row.locked_at = job.lockedAt;
+  if (job.lockedBy !== undefined) row.locked_by = job.lockedBy;
+  if (job.startedAt !== undefined) row.started_at = job.startedAt;
+  if (job.completedAt !== undefined) row.completed_at = job.completedAt;
+  if (job.webhookUrl !== undefined) row.webhook_url = job.webhookUrl;
+  if (job.webhookAttempts !== undefined) row.webhook_attempts = job.webhookAttempts;
+  if (job.webhookLastStatus !== undefined) row.webhook_last_status = job.webhookLastStatus;
  if (job.createdAt !== undefined) row.created_at = job.createdAt;
  if (job.updatedAt !== undefined) row.updated_at = job.updatedAt;
  return row;
@@ -310,6 +470,7 @@ function jobFromRow(row: Record<string, unknown>): GenerationJob {
  return {
    id: String(row.id),
    ownerId: String(row.owner_id),
+    externalClientId: optionalString(row.external_client_id),
    capability: row.capability as GenerationJob["capability"],
    provider: row.provider as GenerationJob["provider"],
    reqKey: String(row.req_key),
@@ -323,6 +484,27 @@ function jobFromRow(row: Record<string, unknown>): GenerationJob {
    responsePayload: isRecord(row.response_payload) ? row.response_payload : undefined,
    error: isRecord(row.error) ? { message: String(row.error.message || "Unknown error"), code: row.error.code as string | number | undefined, retryable: Boolean(row.error.retryable) } : undefined,
    retryOf: row.retry_of ? String(row.retry_of) : undefined,
+    idempotencyKey: optionalString(row.idempotency_key),
+    idempotencyFingerprint: optionalString(row.idempotency_fingerprint),
+    priority: optionalNumber(row.priority),
+    attempts: optionalNumber(row.attempts),
+    maxAttempts: optionalNumber(row.max_attempts),
+    scheduledAt: optionalString(row.scheduled_at),
+    lockedAt: optionalString(row.locked_at),
+    lockedBy: optionalString(row.locked_by),
+    startedAt: optionalString(row.started_at),
+    completedAt: optionalString(row.completed_at),
+    webhookUrl: optionalString(row.webhook_url),
+    webhookAttempts: optionalNumber(row.webhook_attempts),
+    webhookLastStatus: isRecord(row.webhook_last_status)
+      ? {
+          ok: Boolean(row.webhook_last_status.ok),
+          status: optionalNumber(row.webhook_last_status.status),
+          error: optionalString(row.webhook_last_status.error),
+          attemptedAt: String(row.webhook_last_status.attemptedAt || row.webhook_last_status.attempted_at || ""),
+          nextAttemptAt: optionalString(row.webhook_last_status.nextAttemptAt || row.webhook_last_status.next_attempt_at)
+        }
+      : undefined,
    createdAt: String(row.created_at),
    updatedAt: String(row.updated_at)
  };
@@ -355,3 +537,15 @@ function usageFromRow(row: Record<string, unknown>): UsageEvent {
 function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null && !Array.isArray(value);
 }
+
+function optionalString(value: unknown): string | undefined {
+  if (typeof value !== "string") return undefined;
+  const trimmed = value.trim();
+  return trimmed || undefined;
+}
+
+function optionalNumber(value: unknown): number | undefined {
+  if (value === undefined || value === null || value === "") return undefined;
+  const parsed = Number(value);
+  return Number.isFinite(parsed) ? parsed : undefined;
+}