NianAIGC/lib/server/generation-service.ts

import {
  buildJimengPayload,
  buildJimengQueryPayload,
  getEnabledImageCapability,
  isRetryableVisualCode
} from "@/lib/jimeng/capabilities";
import {
  buildEvolinkImagePayload,
  extractEvolinkResultUrls,
  getEffectiveImageEngine,
  getEvolinkImageSettings,
  getEvolinkTaskId,
  mapEvolinkStatus,
  queryEvolinkTask,
  shouldMockEvolinkApi,
  submitEvolinkImageTask,
  type EvolinkTaskResponse
} from "@/lib/evolink/image-client";
import {
  createGenerationJob,
  getGenerationJob,
  recordUsageEvent,
  updateGenerationJob
} from "@/lib/server/data-store";
import { createMockImageBuffer } from "@/lib/server/mock-image";
import { importRemoteImageAsAsset, saveGeneratedAsset } from "@/lib/server/storage";
import { DEFAULT_OWNER_ID, toAbsoluteUrl } from "@/lib/server/runtime";
import type { EnabledImageCapability, GenerationJob, VisualTaskQueryResponse } from "@/lib/types";
import { queryVisualTask, shouldMockVisualApi, submitVisualTask } from "@/lib/volcengine/visual-client";

export type SubmitImageJobInput = {
  ownerId?: string;
  externalClientId?: string;
  capability: EnabledImageCapability;
  prompt?: string;
  imageUrls?: string[];
  inputAssetIds?: string[];
  scale?: number;
  width?: number;
  height?: number;
  min_ratio?: number;
  max_ratio?: number;
  force_single?: boolean;
  resolution?: "4k" | "8k";
  seed?: number;
  retryOf?: string;
  idempotencyKey?: string;
  idempotencyFingerprint?: string;
  priority?: number;
  maxAttempts?: number;
  webhookUrl?: string;
};

export async function submitImageJob(input: SubmitImageJobInput, origin: string): Promise<GenerationJob> {
  const ownerId = input.ownerId || DEFAULT_OWNER_ID;
  const capability = getEnabledImageCapability(input.capability);
  const normalizedUrls = (input.imageUrls || []).map((url) => toAbsoluteUrl(url, origin));
  const engine = getEffectiveImageEngine(input.capability);
  const providerPayload = engine === "evolink"
    ? buildEvolinkImagePayload(input.capability, { ...input, imageUrls: normalizedUrls })
    : buildJimengPayload(input.capability, capability.reqKey, { ...input, imageUrls: normalizedUrls });
  const mock = engine === "evolink" ? shouldMockEvolinkApi() : shouldMockVisualApi();
  const reqKey = engine === "evolink" ? getEvolinkImageSettings().model : capability.reqKey;
  let job = await createGenerationJob({
    ownerId,
    externalClientId: input.externalClientId,
    capability: input.capability,
    provider: mock ? "mock" : engine === "evolink" ? "evolink" : "volcengine-visual",
    reqKey,
    status: "queued",
    prompt: input.prompt,
    inputAssetIds: input.inputAssetIds || [],
    inputUrls: normalizedUrls,
    outputAssetIds: [],
    requestPayload: {
      engine,
      input,
      providerPayload
    },
    retryOf: input.retryOf,
    idempotencyKey: input.idempotencyKey,
    idempotencyFingerprint: input.idempotencyFingerprint,
    priority: input.priority,
    maxAttempts: input.maxAttempts,
    webhookUrl: input.webhookUrl
  });

  return job;
}

export async function advanceImageJob(jobId: string, origin: string): Promise<GenerationJob> {
  const job = await getGenerationJob(jobId);
  if (!job) throw new Error(`Generation job not found: ${jobId}`);
  if (["succeeded", "failed", "expired", "cancelled"].includes(job.status)) return job;
  if (job.provider === "mock") return completeMockJob(job, origin);
  if (!job.providerTaskId) return dispatchImageJob(job);
  return syncImageJob(job.id, origin);
}

async function dispatchImageJob(job: GenerationJob): Promise<GenerationJob> {
  const providerPayload = asRecord(job.requestPayload.providerPayload);
  try {
    if (job.provider === "evolink") {
      const response = await submitEvolinkImageTask(providerPayload);
      const taskId = getEvolinkTaskId(response);
      if (!taskId) {
        job = await updateGenerationJob(job.id, {
          status: "failed",
          responsePayload: response as unknown as Record<string, unknown>,
          error: {
            code: response.error?.code,
            message: response.error?.message || response.message || "EvoLink task submission failed.",
            retryable: false
          }
        });
        return job;
      }
      return updateGenerationJob(job.id, {
        status: "queued",
        providerTaskId: taskId,
        responsePayload: response as unknown as Record<string, unknown>
      });
    }

    const response = await submitVisualTask(providerPayload);
    if (response.code !== 10000 || !response.data?.task_id) {
      job = await updateGenerationJob(job.id, {
        status: "failed",
        responsePayload: response as unknown as Record<string, unknown>,
        error: {
          code: response.code,
          message: response.message || "Volcengine Visual task submission failed.",
          retryable: isRetryableVisualCode(response.code)
        }
      });
      return job;
    }
    return updateGenerationJob(job.id, {
      status: "queued",
      providerTaskId: response.data.task_id,
      responsePayload: response as unknown as Record<string, unknown>
    });
  } catch (error) {
    return updateGenerationJob(job.id, {
      status: "failed",
      error: {
        message: error instanceof Error ? error.message : String(error),
        retryable: false
      }
    });
  }
}

export async function syncImageJob(jobId: string, origin: string): Promise<GenerationJob> {
  const job = await getGenerationJob(jobId);
  if (!job) throw new Error(`Generation job not found: ${jobId}`);
  if (["succeeded", "failed", "expired", "cancelled"].includes(job.status)) return job;
  if (job.provider === "mock") return completeMockJob(job, origin);
  if (!job.providerTaskId) return job;

  if (job.provider === "evolink") {
    return syncEvolinkImageJob(job, origin);
  }

  const queryPayload = buildJimengQueryPayload(job.reqKey, job.providerTaskId);
  let response: VisualTaskQueryResponse;
  try {
    response = await queryVisualTask(queryPayload);
  } catch (error) {
    return updateGenerationJob(job.id, {
      status: "failed",
      error: {
        message: error instanceof Error ? error.message : String(error),
        retryable: false
      }
    });
  }

  if (response.data?.status === "in_queue") {
    return updateGenerationJob(job.id, { status: "queued", responsePayload: response as unknown as Record<string, unknown> });
  }
  if (response.data?.status === "generating") {
    return updateGenerationJob(job.id, { status: "running", responsePayload: response as unknown as Record<string, unknown> });
  }
  if (response.data?.status === "expired" || response.data?.status === "not_found") {
    return updateGenerationJob(job.id, {
      status: "expired",
      responsePayload: response as unknown as Record<string, unknown>,
      error: {
        code: response.code,
        message: response.message || "Volcengine Visual task expired or was not found.",
        retryable: true
      }
    });
  }
  if (response.code !== 10000) {
    return updateGenerationJob(job.id, {
      status: "failed",
      responsePayload: response as unknown as Record<string, unknown>,
      error: {
        code: response.code,
        message: response.message || "Volcengine Visual task failed.",
        retryable: isRetryableVisualCode(response.code)
      }
    });
  }

  const urls = response.data?.image_urls || [];
  const assets = [];
  for (let index = 0; index < urls.length; index += 1) {
    assets.push(await importRemoteImageAsAsset({
      ownerId: job.ownerId,
      url: urls[index],
      origin,
      source: sourceForCapability(job.capability),
      capability: job.capability,
      jobId: job.id,
      index,
      tags: assetTagsForJob(job)
    }));
  }
  if (assets.length) {
    await recordUsageEvent({
      ownerId: job.ownerId,
      jobId: job.id,
      capability: job.capability,
      quantity: assets.length,
      estimatedUnit: "image"
    });
  }
  return updateGenerationJob(job.id, {
    status: "succeeded",
    outputAssetIds: assets.map((asset) => asset.id),
    responsePayload: response as unknown as Record<string, unknown>
  });
}

async function syncEvolinkImageJob(job: GenerationJob, origin: string): Promise<GenerationJob> {
  if (!job.providerTaskId) return job;
  let response: EvolinkTaskResponse;
  try {
    response = await queryEvolinkTask(job.providerTaskId);
  } catch (error) {
    return updateGenerationJob(job.id, {
      status: "failed",
      error: {
        message: error instanceof Error ? error.message : String(error),
        retryable: false
      }
    });
  }

  const status = mapEvolinkStatus(response);
  if (status === "queued" || status === "running") {
    return updateGenerationJob(job.id, {
      status,
      responsePayload: response as unknown as Record<string, unknown>
    });
  }
  if (status === "expired" || status === "cancelled") {
    return updateGenerationJob(job.id, {
      status,
      responsePayload: response as unknown as Record<string, unknown>,
      error: {
        code: response.error?.code,
        message: response.error?.message || response.message || "EvoLink image task was not completed.",
        retryable: status === "expired"
      }
    });
  }
  if (status === "failed") {
    return updateGenerationJob(job.id, {
      status: "failed",
      responsePayload: response as unknown as Record<string, unknown>,
      error: {
        code: response.error?.code,
        message: response.error?.message || response.message || "EvoLink image task failed.",
        retryable: false
      }
    });
  }

  const urls = extractEvolinkResultUrls(response);
  if (!urls.length) {
    return updateGenerationJob(job.id, {
      status: "failed",
      responsePayload: response as unknown as Record<string, unknown>,
      error: {
        message: "EvoLink image task completed without result URLs.",
        retryable: true
      }
    });
  }

  const assets = [];
  for (let index = 0; index < urls.length; index += 1) {
    assets.push(await importRemoteImageAsAsset({
      ownerId: job.ownerId,
      url: urls[index],
      origin,
      source: sourceForCapability(job.capability),
      capability: job.capability,
      jobId: job.id,
      index,
      tags: assetTagsForJob(job)
    }));
  }
  await recordUsageEvent({
    ownerId: job.ownerId,
    jobId: job.id,
    capability: job.capability,
    quantity: assets.length,
    estimatedUnit: "image"
  });
  return updateGenerationJob(job.id, {
    status: "succeeded",
    outputAssetIds: assets.map((asset) => asset.id),
    responsePayload: response as unknown as Record<string, unknown>
  });
}

export async function retryImageJob(jobId: string, origin: string): Promise<GenerationJob> {
  const job = await getGenerationJob(jobId);
  if (!job) throw new Error(`Generation job not found: ${jobId}`);
  const input = (job.requestPayload.input || {}) as SubmitImageJobInput;
  return submitImageJob({
    ...input,
    capability: job.capability as EnabledImageCapability,
    retryOf: job.id
  }, origin);
}

async function completeMockJob(job: GenerationJob, origin: string): Promise<GenerationJob> {
  if (job.status === "succeeded" && job.outputAssetIds.length > 0) return job;
  const requestInput = (job.requestPayload.input || {}) as SubmitImageJobInput;
  const buffer = createMockImageBuffer({
    title: "智念AIGC生成结果",
    prompt: job.prompt || requestInput.prompt,
    capability: job.capability,
    resolution: requestInput.resolution
  });
  const asset = await saveGeneratedAsset({
    ownerId: job.ownerId,
    bytes: buffer,
    fileName: `${job.capability.replace(/\W+/g, "-")}-${job.id}.svg`,
    contentType: "image/svg+xml",
    origin,
    source: sourceForCapability(job.capability),
    capability: job.capability,
    jobId: job.id,
    tags: assetTagsForJob(job),
    metadata: {
      mock: true
    }
  });
  await recordUsageEvent({
    ownerId: job.ownerId,
    jobId: job.id,
    capability: job.capability,
    quantity: 1,
    estimatedUnit: "image"
  });
  return updateGenerationJob(job.id, {
    status: "succeeded",
    outputAssetIds: [asset.id],
    responsePayload: {
      mock: true,
      data: {
        status: "done",
        image_urls: [asset.url]
      }
    }
  });
}

function sourceForCapability(capability: string) {
  if (capability === "image.inpaint") return "edited";
  if (capability === "image.upscale") return "upscaled";
  return "generated";
}

function assetTagsForJob(job: GenerationJob): string[] {
  return job.externalClientId ? [job.capability, `api-client:${job.externalClientId}`] : [job.capability];
}

function asRecord(value: unknown): Record<string, unknown> {
  return typeof value === "object" && value !== null && !Array.isArray(value)
    ? value as Record<string, unknown>
    : {};
}