stack/apps/orchestrator/src/agent-ingestion/agent-ingestion.service.ts

import { Injectable, Logger } from "@nestjs/common";
import type { Prisma } from "@prisma/client";
import { PrismaService } from "../prisma/prisma.service";

export type AgentConversationRole = "agent" | "user" | "system" | "operator";

@Injectable()
export class AgentIngestionService {
  private readonly logger = new Logger(AgentIngestionService.name);

  constructor(private readonly prisma: PrismaService) {}

  private toJsonValue(value: Record<string, unknown>): Prisma.InputJsonValue {
    return value as Prisma.InputJsonValue;
  }

  async recordAgentSpawned(
    agentId: string,
    parentAgentId?: string,
    missionId?: string,
    taskId?: string,
    agentType?: string
  ): Promise<void> {
    await this.prisma.agentSessionTree.upsert({
      where: { sessionId: agentId },
      create: {
        sessionId: agentId,
        parentSessionId: parentAgentId,
        missionId,
        taskId,
        agentType,
        status: "spawning",
      },
      update: {
        parentSessionId: parentAgentId,
        missionId,
        taskId,
        agentType,
        status: "spawning",
        completedAt: null,
      },
    });

    this.logger.debug(`Recorded spawned state for agent ${agentId}`);
  }

  async recordAgentStarted(agentId: string): Promise<void> {
    await this.prisma.agentSessionTree.upsert({
      where: { sessionId: agentId },
      create: {
        sessionId: agentId,
        status: "running",
      },
      update: {
        status: "running",
      },
    });

    this.logger.debug(`Recorded running state for agent ${agentId}`);
  }

  async recordAgentCompleted(agentId: string): Promise<void> {
    const completedAt = new Date();

    await this.prisma.agentSessionTree.upsert({
      where: { sessionId: agentId },
      create: {
        sessionId: agentId,
        status: "completed",
        completedAt,
      },
      update: {
        status: "completed",
        completedAt,
      },
    });

    this.logger.debug(`Recorded completed state for agent ${agentId}`);
  }

  async recordAgentFailed(agentId: string, error?: string): Promise<void> {
    const completedAt = new Date();
    const metadata = error ? this.toJsonValue({ error }) : undefined;

    await this.prisma.agentSessionTree.upsert({
      where: { sessionId: agentId },
      create: {
        sessionId: agentId,
        status: "failed",
        completedAt,
        ...(metadata && { metadata }),
      },
      update: {
        status: "failed",
        completedAt,
        ...(metadata && { metadata }),
      },
    });

    this.logger.debug(`Recorded failed state for agent ${agentId}`);
  }

  async recordAgentKilled(agentId: string): Promise<void> {
    const completedAt = new Date();

    await this.prisma.agentSessionTree.upsert({
      where: { sessionId: agentId },
      create: {
        sessionId: agentId,
        status: "killed",
        completedAt,
      },
      update: {
        status: "killed",
        completedAt,
      },
    });

    this.logger.debug(`Recorded killed state for agent ${agentId}`);
  }

  async recordMessage(
    sessionId: string,
    role: AgentConversationRole,
    content: string,
    provider = "internal",
    metadata?: Record<string, unknown>
  ): Promise<void> {
    await this.prisma.agentConversationMessage.create({
      data: {
        sessionId,
        role,
        content,
        provider,
        ...(metadata && { metadata: this.toJsonValue(metadata) }),
      },
    });

    this.logger.debug(`Recorded message for session ${sessionId}`);
  }
}