stack/apps/api/src/agent-tasks/agent-tasks.service.ts

import { Injectable, NotFoundException } from "@nestjs/common";
import { PrismaService } from "../prisma/prisma.service";
import { AgentTaskStatus, AgentTaskPriority, Prisma } from "@prisma/client";
import type { CreateAgentTaskDto, UpdateAgentTaskDto, QueryAgentTasksDto } from "./dto";

/**
 * Service for managing agent tasks
 */
@Injectable()
export class AgentTasksService {
  constructor(private readonly prisma: PrismaService) {}

  /**
   * Create a new agent task
   */
  async create(workspaceId: string, userId: string, createAgentTaskDto: CreateAgentTaskDto) {
    // Build the create input, handling optional fields properly for exactOptionalPropertyTypes
    const createInput: Prisma.AgentTaskUncheckedCreateInput = {
      title: createAgentTaskDto.title,
      workspaceId,
      createdById: userId,
      status: createAgentTaskDto.status ?? AgentTaskStatus.PENDING,
      priority: createAgentTaskDto.priority ?? AgentTaskPriority.MEDIUM,
      agentType: createAgentTaskDto.agentType,
      agentConfig: (createAgentTaskDto.agentConfig ?? {}) as Prisma.InputJsonValue,
    };

    // Add optional fields only if they exist
    if (createAgentTaskDto.description) createInput.description = createAgentTaskDto.description;
    if (createAgentTaskDto.result)
      createInput.result = createAgentTaskDto.result as Prisma.InputJsonValue;
    if (createAgentTaskDto.error) createInput.error = createAgentTaskDto.error;

    // Set startedAt if status is RUNNING
    if (createInput.status === AgentTaskStatus.RUNNING) {
      createInput.startedAt = new Date();
    }

    // Set completedAt if status is COMPLETED or FAILED
    if (
      createInput.status === AgentTaskStatus.COMPLETED ||
      createInput.status === AgentTaskStatus.FAILED
    ) {
      createInput.completedAt = new Date();
      createInput.startedAt ??= new Date();
    }

    const agentTask = await this.prisma.agentTask.create({
      data: createInput,
      include: {
        createdBy: {
          select: { id: true, name: true, email: true },
        },
      },
    });

    return agentTask;
  }

  /**
   * Get paginated agent tasks with filters
   */
  async findAll(query: QueryAgentTasksDto) {
    const page = query.page ?? 1;
    const limit = query.limit ?? 50;
    const skip = (page - 1) * limit;

    // Build where clause
    const where: Prisma.AgentTaskWhereInput = {};

    if (query.workspaceId) {
      where.workspaceId = query.workspaceId;
    }

    if (query.status) {
      where.status = query.status;
    }

    if (query.priority) {
      where.priority = query.priority;
    }

    if (query.agentType) {
      where.agentType = query.agentType;
    }

    if (query.createdById) {
      where.createdById = query.createdById;
    }

    // Execute queries in parallel
    const [data, total] = await Promise.all([
      this.prisma.agentTask.findMany({
        where,
        include: {
          createdBy: {
            select: { id: true, name: true, email: true },
          },
        },
        orderBy: {
          createdAt: "desc",
        },
        skip,
        take: limit,
      }),
      this.prisma.agentTask.count({ where }),
    ]);

    return {
      data,
      meta: {
        total,
        page,
        limit,
        totalPages: Math.ceil(total / limit),
      },
    };
  }

  /**
   * Get a single agent task by ID
   */
  async findOne(id: string, workspaceId: string) {
    const agentTask = await this.prisma.agentTask.findUnique({
      where: {
        id,
        workspaceId,
      },
      include: {
        createdBy: {
          select: { id: true, name: true, email: true },
        },
      },
    });

    if (!agentTask) {
      throw new NotFoundException(`Agent task with ID ${id} not found`);
    }

    return agentTask;
  }

  /**
   * Update an agent task
   */
  async update(id: string, workspaceId: string, updateAgentTaskDto: UpdateAgentTaskDto) {
    // Verify agent task exists
    const existingTask = await this.prisma.agentTask.findUnique({
      where: { id, workspaceId },
    });

    if (!existingTask) {
      throw new NotFoundException(`Agent task with ID ${id} not found`);
    }

    const data: Prisma.AgentTaskUpdateInput = {};

    // Only include fields that are actually being updated
    if (updateAgentTaskDto.title !== undefined) data.title = updateAgentTaskDto.title;
    if (updateAgentTaskDto.description !== undefined)
      data.description = updateAgentTaskDto.description;
    if (updateAgentTaskDto.status !== undefined) data.status = updateAgentTaskDto.status;
    if (updateAgentTaskDto.priority !== undefined) data.priority = updateAgentTaskDto.priority;
    if (updateAgentTaskDto.agentType !== undefined) data.agentType = updateAgentTaskDto.agentType;
    if (updateAgentTaskDto.error !== undefined) data.error = updateAgentTaskDto.error;

    if (updateAgentTaskDto.agentConfig !== undefined) {
      data.agentConfig = updateAgentTaskDto.agentConfig as Prisma.InputJsonValue;
    }

    if (updateAgentTaskDto.result !== undefined) {
      data.result =
        updateAgentTaskDto.result === null
          ? Prisma.JsonNull
          : (updateAgentTaskDto.result as Prisma.InputJsonValue);
    }

    // Handle startedAt based on status changes
    if (updateAgentTaskDto.status) {
      if (
        updateAgentTaskDto.status === AgentTaskStatus.RUNNING &&
        existingTask.status === AgentTaskStatus.PENDING &&
        !existingTask.startedAt
      ) {
        data.startedAt = new Date();
      }

      // Handle completedAt based on status changes
      if (
        (updateAgentTaskDto.status === AgentTaskStatus.COMPLETED ||
          updateAgentTaskDto.status === AgentTaskStatus.FAILED) &&
        existingTask.status !== AgentTaskStatus.COMPLETED &&
        existingTask.status !== AgentTaskStatus.FAILED
      ) {
        data.completedAt = new Date();
        if (!existingTask.startedAt) {
          data.startedAt = new Date();
        }
      }
    }

    const agentTask = await this.prisma.agentTask.update({
      where: {
        id,
        workspaceId,
      },
      data,
      include: {
        createdBy: {
          select: { id: true, name: true, email: true },
        },
      },
    });

    return agentTask;
  }

  /**
   * Delete an agent task
   */
  async remove(id: string, workspaceId: string) {
    // Verify agent task exists
    const agentTask = await this.prisma.agentTask.findUnique({
      where: { id, workspaceId },
    });

    if (!agentTask) {
      throw new NotFoundException(`Agent task with ID ${id} not found`);
    }

    await this.prisma.agentTask.delete({
      where: {
        id,
        workspaceId,
      },
    });

    return { message: "Agent task deleted successfully" };
  }
}