stack/apps/orchestrator/src/api/agents/agent-tree.service.ts

import { Injectable } from "@nestjs/common";
import type { AgentSessionTree, Prisma } from "@prisma/client";
import { AgentTreeResponseDto } from "./dto/agent-tree-response.dto";
import { PrismaService } from "../../prisma/prisma.service";

const DEFAULT_PAGE_LIMIT = 50;
const MAX_PAGE_LIMIT = 200;

interface SessionCursor {
  spawnedAt: Date;
  sessionId: string;
}

export interface AgentSessionTreeListResult {
  sessions: AgentSessionTree[];
  total: number;
  cursor?: string;
}

@Injectable()
export class AgentTreeService {
  constructor(private readonly prisma: PrismaService) {}

  async listSessions(
    cursor?: string,
    limit = DEFAULT_PAGE_LIMIT
  ): Promise<AgentSessionTreeListResult> {
    const safeLimit = this.normalizeLimit(limit);
    const parsedCursor = this.parseCursor(cursor);

    const where: Prisma.AgentSessionTreeWhereInput | undefined = parsedCursor
      ? {
          OR: [
            {
              spawnedAt: {
                lt: parsedCursor.spawnedAt,
              },
            },
            {
              spawnedAt: parsedCursor.spawnedAt,
              sessionId: {
                lt: parsedCursor.sessionId,
              },
            },
          ],
        }
      : undefined;

    const [sessions, total] = await Promise.all([
      this.prisma.agentSessionTree.findMany({
        where,
        orderBy: [{ spawnedAt: "desc" }, { sessionId: "desc" }],
        take: safeLimit,
      }),
      this.prisma.agentSessionTree.count(),
    ]);

    const nextCursor =
      sessions.length === safeLimit
        ? this.serializeCursor(sessions[sessions.length - 1])
        : undefined;

    return {
      sessions,
      total,
      ...(nextCursor !== undefined ? { cursor: nextCursor } : {}),
    };
  }

  async getSession(sessionId: string): Promise<AgentSessionTree | null> {
    return this.prisma.agentSessionTree.findUnique({
      where: { sessionId },
    });
  }

  async getTree(): Promise<AgentTreeResponseDto[]> {
    const entries = await this.prisma.agentSessionTree.findMany({
      orderBy: { spawnedAt: "desc" },
      take: 200,
    });

    const response: AgentTreeResponseDto[] = [];
    for (const entry of entries) {
      response.push({
        sessionId: entry.sessionId,
        parentSessionId: entry.parentSessionId ?? null,
        status: entry.status,
        agentType: entry.agentType ?? null,
        taskSource: entry.taskSource ?? null,
        spawnedAt: entry.spawnedAt.toISOString(),
        completedAt: entry.completedAt?.toISOString() ?? null,
      });
    }

    return response;
  }

  private normalizeLimit(limit: number): number {
    const normalized = Number.isFinite(limit) ? Math.trunc(limit) : DEFAULT_PAGE_LIMIT;
    if (normalized < 1) {
      return 1;
    }

    return Math.min(normalized, MAX_PAGE_LIMIT);
  }

  private serializeCursor(entry: Pick<AgentSessionTree, "spawnedAt" | "sessionId">): string {
    return Buffer.from(
      JSON.stringify({
        spawnedAt: entry.spawnedAt.toISOString(),
        sessionId: entry.sessionId,
      }),
      "utf8"
    ).toString("base64url");
  }

  private parseCursor(cursor?: string): SessionCursor | null {
    if (!cursor) {
      return null;
    }

    try {
      const decoded = Buffer.from(cursor, "base64url").toString("utf8");
      const parsed = JSON.parse(decoded) as {
        spawnedAt?: string;
        sessionId?: string;
      };

      if (typeof parsed.spawnedAt !== "string" || typeof parsed.sessionId !== "string") {
        return null;
      }

      const spawnedAt = new Date(parsed.spawnedAt);
      if (Number.isNaN(spawnedAt.getTime())) {
        return null;
      }

      return {
        spawnedAt,
        sessionId: parsed.sessionId,
      };
    } catch {
      return null;
    }
  }
}