api-server/src/modules/ai-job/ai-job-lifecycle.repository.ts

import { Injectable } from '@nestjs/common';
import { PrismaService } from '../../infrastructure/database/prisma.service';
import { Prisma } from '@prisma/client';
import {
  JobLockConflictError,
  JobNotActiveError,
  JobAlreadyTerminalError,
  JobNotCancellableError,
  JobExtendLockError,
} from './ai-job.errors';
import {
  AiJobStateMachine,
  LifecycleStatus,
  LIFECYCLE_TO_STATUS,
} from './ai-job-state-machine';

/**
 * ── Job 创建输入 ──
 */
export interface CreateJobInput {
  userId: string;
  jobType: string;
  triggerType?: string;
  queueName?: string;
  targetType?: string;
  targetId?: string;
  parentJobId?: string;
  idempotencyKey?: string;
  priority?: number;
  credentialMode?: string;
  credentialId?: string;
  modelProvider?: string;
  modelName?: string;
  promptKey?: string;
  promptVersion?: string;
  outputSchemaVersion?: string;
  inputSchemaVersion?: string;
  inputRef?: string;
  maxAttempts?: number;
  timeoutMs?: number;
}

/**
 * ── Job 错误信息 ──
 */
export interface JobErrorInfo {
  errorCode: string;
  publicErrorMessage?: string;
  internalErrorMessage?: string;
}

/**
 * M-AI-03-03: AiJobLifecycleRepository
 *
 * 统一 lifecycleStatus 的唯一写入入口。
 *
 * 约束：
 * - 所有状态变更使用 CAS（条件 UPDATE WHERE lifecycleStatus IN (...)）
 * - 禁止先读后无条件 update
 * - 可选 tx 参数支持在外部 Prisma Transaction 中复用事务
 * - lockJob 不使用 tx（CAS 必须立即对其他 Worker 可见）
 */
@Injectable()
export class AiJobLifecycleRepository {
  private readonly LOCK_TTL_MS = 60_000; // ADR-003：60s

  constructor(
    private readonly prisma: PrismaService,
    private readonly stateMachine: AiJobStateMachine,
  ) {}

  // ═══════════════════════════════════════════════════════════
  // 创建
  // ═══════════════════════════════════════════════════════════

  /**
   * 在给定事务中创建 Job（lifecycleStatus = 'queued'）。
   * 调用方负责在同一 tx 中写入 OutboxEvent。
   */
  async createJob(
    tx: Prisma.TransactionClient,
    input: CreateJobInput,
  ) {
    const job = await tx.aiJob.create({
      data: {
        userId: input.userId,
        jobType: input.jobType,
        triggerType: input.triggerType ?? 'user_api',
        queueName: input.queueName ?? 'ai-interactive',
        targetType: input.targetType ?? null,
        targetId: input.targetId ?? null,
        parentJobId: input.parentJobId ?? null,
        idempotencyKey: input.idempotencyKey ?? null,
        priority: input.priority ?? 0,
        credentialMode: input.credentialMode ?? 'platform_key',
        credentialId: input.credentialId ?? null,
        // modelProvider / modelName 不在此层设置默认值。
        // 调用方（AiJobService）负责从 JobDefinition 解析后传入。
        // 若未传入，依赖 Prisma schema @default 保证列非空。
        ...(input.modelProvider ? { modelProvider: input.modelProvider } : {}),
        ...(input.modelName ? { modelName: input.modelName } : {}),
        promptKey: input.promptKey ?? null,
        promptVersion: input.promptVersion ?? null,
        outputSchemaVersion: input.outputSchemaVersion ?? null,
        inputSchemaVersion: input.inputSchemaVersion ?? null,
        inputRef: input.inputRef ?? null,
        maxAttempts: input.maxAttempts ?? 3,
        timeoutMs: input.timeoutMs ?? 120000,
        // ── 生命周期 ──
        lifecycleStatus: 'queued',
        status: LIFECYCLE_TO_STATUS['queued'], // Shadow Write: pending
        queuedAt: new Date(),
        attemptCount: 0,
        progress: 0,
      },
    });

    return job;
  }

  // ═══════════════════════════════════════════════════════════
  // 抢锁（CAS: queued → running）
  // ═══════════════════════════════════════════════════════════

  /**
   * CAS 抢锁。只有 lifecycleStatus = 'queued' 的 Job 才能被锁定。
   *
   * 必须是独立操作（不传 tx）—— CAS 需要立即对其他 Worker 可见。
   */
  async lockJob(jobId: string, instanceId: string) {
    const now = new Date();
    const lockUntil = new Date(now.getTime() + this.LOCK_TTL_MS);

    // CAS: queued → running，仅当 status 仍为 queued 时成功
    const result = await this.prisma.aiJob.updateMany({
      where: {
        id: jobId,
        lifecycleStatus: 'queued',
      },
      data: {
        lifecycleStatus: 'running',
        status: LIFECYCLE_TO_STATUS['running'], // Shadow Write: processing
        lockedBy: instanceId,
        lockedAt: now,
        lockUntil,
        startedAt: now,
        attemptCount: { increment: 1 }, // 原子递增，每次 lockJob 真实执行次数 +1
      },
    });

    if (result.count === 0) {
      // 抢锁失败 —— 检查原因以提供更精确的错误
      const current = await this.prisma.aiJob.findUnique({
        where: { id: jobId },
        select: { lifecycleStatus: true },
      });
      if (!current) {
        throw new JobLockConflictError(jobId); // Job 不存在
      }
      if (this.stateMachine.isTerminal(current.lifecycleStatus as LifecycleStatus)) {
        throw new JobAlreadyTerminalError(jobId, current.lifecycleStatus!);
      }
      if (current.lifecycleStatus === 'running') {
        throw new JobLockConflictError(jobId); // 已被其他 Worker 抢走
      }
      throw new JobLockConflictError(jobId);
    }

    const job = await this.prisma.aiJob.findUniqueOrThrow({
      where: { id: jobId },
    });

    return job;
  }

  // ═══════════════════════════════════════════════════════════
  // 续锁（Heartbeat）
  // ═══════════════════════════════════════════════════════════

  /**
   * 续约锁定期。仅当 Job 处于 'running' 状态且 lockedBy 匹配时成功。
   */
  async extendLock(jobId: string, instanceId: string) {
    const now = new Date();
    const lockUntil = new Date(now.getTime() + this.LOCK_TTL_MS);

    const result = await this.prisma.aiJob.updateMany({
      where: {
        id: jobId,
        lifecycleStatus: 'running',
        lockedBy: instanceId,
      },
      data: { lockUntil },
    });

    if (result.count === 0) {
      const current = await this.prisma.aiJob.findUnique({
        where: { id: jobId },
        select: { lifecycleStatus: true, lockedBy: true },
      });
      if (!current || current.lifecycleStatus !== 'running') {
        throw new JobNotActiveError(jobId, current?.lifecycleStatus ?? 'unknown');
      }
      throw new JobExtendLockError(jobId, instanceId);
    }

    // 读取当前状态以检查 cancelRequestedAt
    const job = await this.prisma.aiJob.findUniqueOrThrow({
      where: { id: jobId },
    });

    return {
      jobId,
      lockUntil: lockUntil.getTime(),
      cancelRequested: job.cancelRequestedAt != null,
    };
  }

  // ═══════════════════════════════════════════════════════════
  // 标记成功
  // ═══════════════════════════════════════════════════════════

  /**
   * 标记 Job 成功。
   *
   * @param tx 可选事务客户端。Projector 路径传入以共享事务；非 Projector 路径不传。
   */
  async markSucceeded(jobId: string, tx?: Prisma.TransactionClient): Promise<any> {
    const client = tx ?? this.prisma;
    const now = new Date();

    const result = await client.aiJob.updateMany({
      where: {
        id: jobId,
        lifecycleStatus: 'running',
      },
      data: {
        lifecycleStatus: 'succeeded',
        status: LIFECYCLE_TO_STATUS['succeeded'], // Shadow Write: completed
        finishedAt: now,
        // 清除运行时字段
        lockedBy: null,
        lockedAt: null,
        lockUntil: null,
        errorCode: null,
        publicErrorMessage: null,
        internalErrorMessage: null,
      },
    });

    if (result.count === 0) {
      const current = await client.aiJob.findUnique({
        where: { id: jobId },
        select: { lifecycleStatus: true },
      });
      if (current && this.stateMachine.isTerminal(current.lifecycleStatus as LifecycleStatus)) {
        throw new JobAlreadyTerminalError(jobId, current.lifecycleStatus!);
      }
      throw new JobNotActiveError(jobId, current?.lifecycleStatus ?? 'unknown');
    }

    return client.aiJob.findUniqueOrThrow({ where: { id: jobId } });
  }

  // ═══════════════════════════════════════════════════════════
  // 标记失败
  // ═══════════════════════════════════════════════════════════

  /**
   * 标记 Job 失败（终态）。
   *
   * @param tx 可选事务客户端。Engine 非 Projector 错误路径不传。
   */
  async markFailed(jobId: string, error: JobErrorInfo, tx?: Prisma.TransactionClient): Promise<any> {
    const client = tx ?? this.prisma;
    const now = new Date();

    const result = await client.aiJob.updateMany({
      where: {
        id: jobId,
        lifecycleStatus: 'running',
      },
      data: {
        lifecycleStatus: 'failed',
        status: LIFECYCLE_TO_STATUS['failed'], // Shadow Write: failed
        finishedAt: now,
        lockedBy: null,
        lockedAt: null,
        lockUntil: null,
        errorCode: error.errorCode,
        publicErrorMessage: error.publicErrorMessage ?? null,
        internalErrorMessage: error.internalErrorMessage ?? null,
      },
    });

    if (result.count === 0) {
      const current = await client.aiJob.findUnique({
        where: { id: jobId },
        select: { lifecycleStatus: true },
      });
      if (current && this.stateMachine.isTerminal(current.lifecycleStatus as LifecycleStatus)) {
        throw new JobAlreadyTerminalError(jobId, current.lifecycleStatus!);
      }
      throw new JobNotActiveError(jobId, current?.lifecycleStatus ?? 'unknown');
    }

    return client.aiJob.findUniqueOrThrow({ where: { id: jobId } });
  }

  // ═══════════════════════════════════════════════════════════
  // 标记取消
  // ═══════════════════════════════════════════════════════════

  /**
   * 标记 Job 取消。
   *
   * queued 状态可直接取消；running 状态需先设置 cancelRequestedAt，
   * Engine/Heartbeat 检测后调用此方法。
   *
   * @param tx 可选事务客户端。取消路径通常不传（独立操作）。
   */
  async markCancelled(jobId: string, tx?: Prisma.TransactionClient): Promise<any> {
    const client = tx ?? this.prisma;
    const now = new Date();

    // 尝试从 queued 或 running 取消
    const result = await client.aiJob.updateMany({
      where: {
        id: jobId,
        lifecycleStatus: { in: ['queued', 'running'] },
      },
      data: {
        lifecycleStatus: 'cancelled',
        // ── Shadow Write：cancelled 在旧 status 中不存在，映射为 failed + errorCode=CANCELLED ──
        status: LIFECYCLE_TO_STATUS['cancelled'], // 'failed'
        finishedAt: now,
        lockedBy: null,
        lockedAt: null,
        lockUntil: null,
        errorCode: 'CANCELLED',
        publicErrorMessage: 'Job was cancelled',
        internalErrorMessage: null,
      },
    });

    if (result.count === 0) {
      const current = await client.aiJob.findUnique({
        where: { id: jobId },
        select: { lifecycleStatus: true },
      });
      if (!current) {
        throw new JobLockConflictError(jobId); // Job 不存在，复用此错误
      }
      if (this.stateMachine.isTerminal(current.lifecycleStatus as LifecycleStatus)) {
        throw new JobNotCancellableError(jobId, current.lifecycleStatus!);
      }
      throw new JobNotCancellableError(jobId, current.lifecycleStatus ?? 'unknown');
    }

    return client.aiJob.findUniqueOrThrow({ where: { id: jobId } });
  }

  // ═══════════════════════════════════════════════════════════
  // 设置取消请求（不改变 lifecycleStatus）
  // ═══════════════════════════════════════════════════════════

  /**
   * 对 running 状态的 Job 设置 cancelRequestedAt 信号。
   * 不改变 lifecycleStatus — Engine/Heartbeat 检测后调用 markCancelled。
   *
   * 对 queued 状态的 Job 直接取消。
   */
  async requestCancellation(jobId: string): Promise<{ status: 'cancelled' | 'cancel_requested' }> {
    const job = await this.prisma.aiJob.findUnique({
      where: { id: jobId },
      select: { lifecycleStatus: true },
    });

    if (!job) throw new JobLockConflictError(jobId);

    const status = this.stateMachine.parse(job.lifecycleStatus);

    if (status === 'queued') {
      // 直接取消（markCancelled 内部已有 CAS + 错误分类）
      await this.markCancelled(jobId);
      return { status: 'cancelled' };
    }

    if (status === 'running') {
      // CAS 写入 cancelRequestedAt — 只有 lifecycleStatus 仍为 running 时才成功。
      // 使用 updateMany 避免 read→write 之间的 TOCTOU 竞态窗口。
      const result = await this.prisma.aiJob.updateMany({
        where: { id: jobId, lifecycleStatus: 'running' },
        data: { cancelRequestedAt: new Date() },
      });

      if (result.count === 0) {
        // 在 findUnique 和 updateMany 之间 Job 已进入终态或被其他 Worker 修改。
        // 重新读取以返回精确错误。
        const current = await this.prisma.aiJob.findUnique({
          where: { id: jobId },
          select: { lifecycleStatus: true },
        });
        const currentStatus = current?.lifecycleStatus ?? 'unknown';
        if (currentStatus === 'cancelled') {
          return { status: 'cancelled' }; // 已被其他路径取消，幂等返回
        }
        if (this.stateMachine.isTerminal(
          this.stateMachine.parse(currentStatus),
        )) {
          throw new JobNotCancellableError(jobId, currentStatus);
        }
        throw new JobNotCancellableError(jobId, currentStatus);
      }

      return { status: 'cancel_requested' };
    }

    // 终态不可取消
    throw new JobNotCancellableError(jobId, status);
  }

  // ═══════════════════════════════════════════════════════════
  // 只读查询
  // ═══════════════════════════════════════════════════════════

  async findByLifecycleStatus(
    statuses: LifecycleStatus[],
    opts?: { queueName?: string; limit?: number },
  ) {
    const where: Prisma.AiJobWhereInput = {
      lifecycleStatus: { in: statuses },
    };
    if (opts?.queueName) where.queueName = opts.queueName;

    return this.prisma.aiJob.findMany({
      where,
      orderBy: [{ priority: 'asc' }, { createdAt: 'asc' }],
      take: opts?.limit ?? 50,
    });
  }

  async findExpiredLocks(thresholdMs: number = 90_000) {
    const cutoff = new Date(Date.now() - thresholdMs);
    return this.prisma.aiJob.findMany({
      where: {
        lifecycleStatus: 'running',
        lockUntil: { lt: cutoff },
      },
      orderBy: { createdAt: 'asc' },
    });
  }

  async findById(jobId: string) {
    return this.prisma.aiJob.findUnique({ where: { id: jobId } });
  }
}