Genarrative/src/components/image-editor/model3d-generation/useModel3dGenerationTask.ts

import { useCallback } from 'react';

import type { ExternalGenerationJobStatusRecord } from '../../../../packages/shared/src/contracts/externalGeneration';
import type {
  EditorCanvasGenerationCompletionInput,
  EditorProjectSnapshot,
} from '../../../services/image-editor/editorProjectClient';
import {
  submitModel3dImageToModelRequest,
  submitModel3dTextToModelRequest,
} from '../../../services/image-editor/editorProjectClient';
import type {
  CanvasGenerationDialogState,
  GenerateDialogState,
} from '../ImageCanvasEditorTypes';
import { readEditorModel3dPricingConfig } from '../ImageCanvasGenerationModel';
import { isModel3dGenerationMode } from './Model3dGenerationFormModel';
import {
  buildModel3dSubmissionPlan,
  MODEL3D_QUEUE_HANDOFF_FAILED_MESSAGE,
  MODEL3D_SUBMISSION_MISSING_QUEUE_STATE_MESSAGE,
  model3dQueueStateFromResponse,
  resolveModel3dGenerationErrorMessage,
  resolveModel3dResultTitle,
} from './Model3dGenerationSubmission';

type ApplyQueuedGeneration = (
  response: { queueState?: ExternalGenerationJobStatusRecord | null },
  projectId: string | null | undefined,
  applyProjectSnapshot: ((project: EditorProjectSnapshot) => void) | undefined,
  onQueuedGenerationTask?: () => void,
  onWalletBalanceMayHaveChanged?: () => void,
  onGenerationWarning?: (message: string) => void,
  completionDialogId?: string | null,
) => Promise<boolean>;

type UpdateCanvasGenerationDialogById = (
  dialogId: string,
  updater: (
    dialog: CanvasGenerationDialogState,
  ) => CanvasGenerationDialogState | null,
) => void;

/**
 * 3D 生成的提交编排：只负责校验、入队、轮询与收口，请求构造在
 * `Model3dGenerationSubmission`，画布与任务侧栏复用既有链路。
 */
export function useModel3dGenerationTask({
  projectId,
  assetFolderId,
  updateCanvasGenerationDialogById,
  applyProjectSnapshot,
  getGeneratingDialogPlaceholder,
  applyQueuedGeneration,
  onQueuedGenerationTask,
  onWalletBalanceMayHaveChanged,
  onGenerationWarning,
}: {
  projectId?: string | null;
  /** 当前素材夹：与其它画布生成工具一致，结果同时落素材库。 */
  assetFolderId?: string | null;
  updateCanvasGenerationDialogById: UpdateCanvasGenerationDialogById;
  applyProjectSnapshot?: (project: EditorProjectSnapshot) => void;
  getGeneratingDialogPlaceholder: (
    dialog: GenerateDialogState,
  ) => EditorCanvasGenerationCompletionInput['placeholder'] | null | undefined;
  applyQueuedGeneration: ApplyQueuedGeneration;
  onQueuedGenerationTask?: () => void;
  onWalletBalanceMayHaveChanged?: () => void;
  onGenerationWarning?: (message: string) => void;
}) {
  const submitModel3dGeneration = useCallback(
    async (dialog: GenerateDialogState) => {
      const dialogId = dialog.id;
      if (!dialogId || !isModel3dGenerationMode(dialog.mode)) {
        return;
      }
      // 整个提交编排都落在 try 里：取占位框、算标题、构造提交计划同样可能抛
      // （注入的回调、老快照缺字段），漏出去就成了一次没有用户提示的未捕获 rejection。
      try {
        const placeholder = getGeneratingDialogPlaceholder(dialog);
        const canvasCompletion =
          projectId && placeholder
            ? {
                dialogId,
                title: resolveModel3dResultTitle(dialog),
                placeholder,
              }
            : null;
        const planResult = buildModel3dSubmissionPlan({
          dialog,
          projectId,
          canvasCompletion,
          assetFolderId,
          assetLabel: resolveModel3dResultTitle(dialog),
          pricing: readEditorModel3dPricingConfig(),
        });
        if (!planResult.ok) {
          updateCanvasGenerationDialogById(dialogId, (currentDialog) => ({
            ...currentDialog,
            status: 'failed',
            composerOpen: true,
            errorMessage: planResult.message,
          }));
          return;
        }
        updateCanvasGenerationDialogById(dialogId, (currentDialog) => ({
          ...currentDialog,
          // 兜底铸的代次要写回状态：否则老快照每次提交都换幂等键，去重保护失效。
          model3dAttemptNonce: planResult.plan.attemptNonce,
          status: 'generating',
          composerOpen: false,
          errorMessage: undefined,
        }));
        const plan = planResult.plan;
        // 判别联合在这里收窄：文生与图生各走自己的客户端函数，端点与提交体的配对由类型保证。
        const response =
          plan.endpoint === 'text-to-model'
            ? await submitModel3dTextToModelRequest({
                body: plan.body,
                idempotencyKey: plan.requestKey,
              })
            : await submitModel3dImageToModelRequest({
                body: plan.body,
                idempotencyKey: plan.requestKey,
              });
        const queueState = model3dQueueStateFromResponse(response);
        if (!queueState) {
          // 2xx 但没有任务状态（空 body 或缺 status 的 payload）不能当成功静默复位：
          // 用户既看不到任务也没有任何提示，服务端细节被丢掉。
          // 服务端在这条路上也可能已经扣过费，钱包余额同样要按「可能变了」刷新。
          onWalletBalanceMayHaveChanged?.();
          updateCanvasGenerationDialogById(dialogId, (currentDialog) => ({
            ...currentDialog,
            status: 'failed',
            composerOpen: true,
            errorMessage: MODEL3D_SUBMISSION_MISSING_QUEUE_STATE_MESSAGE,
          }));
          return;
        }
        if (
          await applyQueuedGeneration(
            { queueState },
            projectId,
            applyProjectSnapshot,
            onQueuedGenerationTask,
            onWalletBalanceMayHaveChanged,
            onGenerationWarning,
            dialogId,
          )
        ) {
          return;
        }
        // 画布没接住这份任务状态时不能静默复位成 idle：用户会以为没提交过，
        // 再点一次就是又一次外部生成（已经扣过费）。按失败收口，让用户先去看结果。
        updateCanvasGenerationDialogById(dialogId, (currentDialog) => ({
          ...currentDialog,
          status: 'failed',
          composerOpen: true,
          errorMessage: MODEL3D_QUEUE_HANDOFF_FAILED_MESSAGE,
        }));
      } catch (error) {
        // 超时等失败也可能意味着服务端已经扣费（与其它生成链路的
        // runEditorGenerationWithWalletRefresh 同一口径），失败路径也要刷新钱包。
        onWalletBalanceMayHaveChanged?.();
        updateCanvasGenerationDialogById(dialogId, (currentDialog) => ({
          ...currentDialog,
          status: 'failed',
          composerOpen: true,
          errorMessage: resolveModel3dGenerationErrorMessage(error),
        }));
      }
    },
    [
      applyProjectSnapshot,
      applyQueuedGeneration,
      assetFolderId,
      getGeneratingDialogPlaceholder,
      onGenerationWarning,
      onQueuedGenerationTask,
      onWalletBalanceMayHaveChanged,
      projectId,
      updateCanvasGenerationDialogById,
    ],
  );

  return { submitModel3dGeneration };
}