stack/apps/api/src/job-steps/job-steps.service.ts

import { Injectable, NotFoundException } from "@nestjs/common";
import { Prisma, JobStepStatus } from "@prisma/client";
import { PrismaService } from "../prisma/prisma.service";
import type { CreateStepDto, UpdateStepDto } from "./dto";

/**
 * Service for managing job steps
 */
@Injectable()
export class JobStepsService {
  constructor(private readonly prisma: PrismaService) {}

  /**
   * Create a new job step
   */
  async create(jobId: string, createStepDto: CreateStepDto) {
    const data: Prisma.JobStepCreateInput = {
      job: { connect: { id: jobId } },
      ordinal: createStepDto.ordinal,
      phase: createStepDto.phase,
      name: createStepDto.name,
      type: createStepDto.type,
      status: createStepDto.status ?? JobStepStatus.PENDING,
    };

    return this.prisma.jobStep.create({ data });
  }

  /**
   * Get all steps for a job, ordered by ordinal
   */
  async findAllByJob(jobId: string) {
    return this.prisma.jobStep.findMany({
      where: { jobId },
      orderBy: { ordinal: "asc" },
    });
  }

  /**
   * Get a single step by ID
   */
  async findOne(id: string, jobId: string) {
    const step = await this.prisma.jobStep.findUnique({
      where: { id, jobId },
    });

    if (!step) {
      throw new NotFoundException(`JobStep with ID ${id} not found`);
    }

    return step;
  }

  /**
   * Update a job step
   */
  async update(id: string, jobId: string, updateStepDto: UpdateStepDto) {
    // Verify step exists
    await this.findOne(id, jobId);

    const data: Prisma.JobStepUpdateInput = {};

    if (updateStepDto.status !== undefined) {
      data.status = updateStepDto.status;
    }
    if (updateStepDto.output !== undefined) {
      data.output = updateStepDto.output;
    }
    if (updateStepDto.tokensInput !== undefined) {
      data.tokensInput = updateStepDto.tokensInput;
    }
    if (updateStepDto.tokensOutput !== undefined) {
      data.tokensOutput = updateStepDto.tokensOutput;
    }

    return this.prisma.jobStep.update({
      where: { id, jobId },
      data,
    });
  }

  /**
   * Mark a step as running and set startedAt timestamp
   */
  async startStep(id: string, jobId: string) {
    // Verify step exists
    await this.findOne(id, jobId);

    return this.prisma.jobStep.update({
      where: { id, jobId },
      data: {
        status: JobStepStatus.RUNNING,
        startedAt: new Date(),
      },
    });
  }

  /**
   * Mark a step as completed, set output, and calculate duration
   */
  async completeStep(id: string, jobId: string, output?: string) {
    // Verify step exists and get startedAt
    const existingStep = await this.findOne(id, jobId);

    const completedAt = new Date();
    const durationMs = existingStep.startedAt
      ? completedAt.getTime() - existingStep.startedAt.getTime()
      : null;

    const data: Prisma.JobStepUpdateInput = {
      status: JobStepStatus.COMPLETED,
      completedAt,
      durationMs,
    };

    if (output !== undefined) {
      data.output = output;
    }

    return this.prisma.jobStep.update({
      where: { id, jobId },
      data,
    });
  }

  /**
   * Mark a step as failed, set error output, and calculate duration
   */
  async failStep(id: string, jobId: string, error: string) {
    // Verify step exists and get startedAt
    const existingStep = await this.findOne(id, jobId);

    const completedAt = new Date();
    const durationMs = existingStep.startedAt
      ? completedAt.getTime() - existingStep.startedAt.getTime()
      : null;

    return this.prisma.jobStep.update({
      where: { id, jobId },
      data: {
        status: JobStepStatus.FAILED,
        output: error,
        completedAt,
        durationMs,
      },
    });
  }

  /**
   * Start a step - simplified API without jobId
   */
  async start(id: string): Promise<Awaited<ReturnType<typeof this.prisma.jobStep.update>>> {
    const step = await this.prisma.jobStep.findUnique({
      where: { id },
    });

    if (!step) {
      throw new NotFoundException(`JobStep with ID ${id} not found`);
    }

    return this.startStep(id, step.jobId);
  }

  /**
   * Complete a step - simplified API without jobId
   */
  async complete(
    id: string,
    data?: { output?: string; tokensInput?: number; tokensOutput?: number }
  ): Promise<Awaited<ReturnType<typeof this.prisma.jobStep.update>>> {
    const step = await this.prisma.jobStep.findUnique({
      where: { id },
    });

    if (!step) {
      throw new NotFoundException(`JobStep with ID ${id} not found`);
    }

    const existingStep = await this.findOne(id, step.jobId);
    const completedAt = new Date();
    const durationMs = existingStep.startedAt
      ? completedAt.getTime() - existingStep.startedAt.getTime()
      : null;

    const updateData: Prisma.JobStepUpdateInput = {
      status: JobStepStatus.COMPLETED,
      completedAt,
      durationMs,
    };

    if (data?.output !== undefined) {
      updateData.output = data.output;
    }
    if (data?.tokensInput !== undefined) {
      updateData.tokensInput = data.tokensInput;
    }
    if (data?.tokensOutput !== undefined) {
      updateData.tokensOutput = data.tokensOutput;
    }

    return this.prisma.jobStep.update({
      where: { id, jobId: step.jobId },
      data: updateData,
    });
  }

  /**
   * Fail a step - simplified API without jobId
   */
  async fail(
    id: string,
    data?: { error?: string }
  ): Promise<Awaited<ReturnType<typeof this.prisma.jobStep.update>>> {
    const step = await this.prisma.jobStep.findUnique({
      where: { id },
    });

    if (!step) {
      throw new NotFoundException(`JobStep with ID ${id} not found`);
    }

    return this.failStep(id, step.jobId, data?.error ?? "Step failed");
  }

  /**
   * Get steps by job - alias for findAllByJob
   */
  async findByJob(jobId: string): Promise<Awaited<ReturnType<typeof this.findAllByJob>>> {
    return this.findAllByJob(jobId);
  }
}