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

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { AgentTreeService } from "./agent-tree.service";
import { PrismaService } from "../../prisma/prisma.service";

describe("AgentTreeService", () => {
  let service: AgentTreeService;
  let prisma: {
    agentSessionTree: {
      findMany: ReturnType<typeof vi.fn>;
      count: ReturnType<typeof vi.fn>;
      findUnique: ReturnType<typeof vi.fn>;
    };
  };

  beforeEach(() => {
    prisma = {
      agentSessionTree: {
        findMany: vi.fn(),
        count: vi.fn(),
        findUnique: vi.fn(),
      },
    };

    service = new AgentTreeService(prisma as unknown as PrismaService);
  });

  afterEach(() => {
    vi.clearAllMocks();
  });

  describe("listSessions", () => {
    it("returns paginated sessions and cursor", async () => {
      const sessions = [
        {
          id: "tree-2",
          sessionId: "agent-2",
          parentSessionId: null,
          provider: "internal",
          missionId: null,
          taskId: "task-2",
          taskSource: "queue",
          agentType: "worker",
          status: "running",
          spawnedAt: new Date("2026-03-07T11:00:00.000Z"),
          completedAt: null,
          metadata: {},
        },
        {
          id: "tree-1",
          sessionId: "agent-1",
          parentSessionId: null,
          provider: "internal",
          missionId: null,
          taskId: "task-1",
          taskSource: "queue",
          agentType: "worker",
          status: "running",
          spawnedAt: new Date("2026-03-07T10:00:00.000Z"),
          completedAt: null,
          metadata: {},
        },
      ];

      prisma.agentSessionTree.findMany.mockResolvedValue(sessions);
      prisma.agentSessionTree.count.mockResolvedValue(7);

      const result = await service.listSessions(undefined, 2);

      expect(prisma.agentSessionTree.findMany).toHaveBeenCalledWith({
        where: undefined,
        orderBy: [{ spawnedAt: "desc" }, { sessionId: "desc" }],
        take: 2,
      });
      expect(prisma.agentSessionTree.count).toHaveBeenCalledWith();
      expect(result.sessions).toEqual(sessions);
      expect(result.total).toBe(7);
      expect(result.cursor).toBeTypeOf("string");
    });

    it("applies cursor filter when provided", async () => {
      prisma.agentSessionTree.findMany.mockResolvedValue([]);
      prisma.agentSessionTree.count.mockResolvedValue(0);

      const cursorDate = "2026-03-07T10:00:00.000Z";
      const cursorSessionId = "agent-5";
      const cursor = Buffer.from(
        JSON.stringify({
          spawnedAt: cursorDate,
          sessionId: cursorSessionId,
        }),
        "utf8"
      ).toString("base64url");

      await service.listSessions(cursor, 25);

      expect(prisma.agentSessionTree.findMany).toHaveBeenCalledWith({
        where: {
          OR: [
            {
              spawnedAt: {
                lt: new Date(cursorDate),
              },
            },
            {
              spawnedAt: new Date(cursorDate),
              sessionId: {
                lt: cursorSessionId,
              },
            },
          ],
        },
        orderBy: [{ spawnedAt: "desc" }, { sessionId: "desc" }],
        take: 25,
      });
    });

    it("ignores invalid cursor values", async () => {
      prisma.agentSessionTree.findMany.mockResolvedValue([]);
      prisma.agentSessionTree.count.mockResolvedValue(0);

      await service.listSessions("invalid-cursor", 10);

      expect(prisma.agentSessionTree.findMany).toHaveBeenCalledWith({
        where: undefined,
        orderBy: [{ spawnedAt: "desc" }, { sessionId: "desc" }],
        take: 10,
      });
    });
  });

  describe("getSession", () => {
    it("returns matching session entry", async () => {
      const session = {
        id: "tree-1",
        sessionId: "agent-123",
        parentSessionId: null,
        provider: "internal",
        missionId: null,
        taskId: "task-1",
        taskSource: "queue",
        agentType: "worker",
        status: "running",
        spawnedAt: new Date("2026-03-07T11:00:00.000Z"),
        completedAt: null,
        metadata: {},
      };
      prisma.agentSessionTree.findUnique.mockResolvedValue(session);

      const result = await service.getSession("agent-123");

      expect(prisma.agentSessionTree.findUnique).toHaveBeenCalledWith({
        where: { sessionId: "agent-123" },
      });
      expect(result).toEqual(session);
    });

    it("returns null when session does not exist", async () => {
      prisma.agentSessionTree.findUnique.mockResolvedValue(null);

      const result = await service.getSession("agent-missing");

      expect(result).toBeNull();
    });
  });

  describe("getTree", () => {
    it("returns mapped entries from Prisma", async () => {
      prisma.agentSessionTree.findMany.mockResolvedValue([
        {
          id: "tree-1",
          sessionId: "agent-1",
          parentSessionId: "agent-root",
          provider: "internal",
          missionId: "mission-1",
          taskId: "task-1",
          taskSource: "queue",
          agentType: "worker",
          status: "running",
          spawnedAt: new Date("2026-03-07T10:00:00.000Z"),
          completedAt: new Date("2026-03-07T11:00:00.000Z"),
          metadata: {},
        },
      ]);

      const result = await service.getTree();

      expect(prisma.agentSessionTree.findMany).toHaveBeenCalledWith({
        orderBy: { spawnedAt: "desc" },
        take: 200,
      });
      expect(result).toEqual([
        {
          sessionId: "agent-1",
          parentSessionId: "agent-root",
          status: "running",
          agentType: "worker",
          taskSource: "queue",
          spawnedAt: "2026-03-07T10:00:00.000Z",
          completedAt: "2026-03-07T11:00:00.000Z",
        },
      ]);
    });

    it("returns empty array when no entries exist", async () => {
      prisma.agentSessionTree.findMany.mockResolvedValue([]);

      const result = await service.getTree();

      expect(result).toEqual([]);
    });

    it("maps null parentSessionId and completedAt correctly", async () => {
      prisma.agentSessionTree.findMany.mockResolvedValue([
        {
          id: "tree-2",
          sessionId: "agent-root",
          parentSessionId: null,
          provider: "internal",
          missionId: null,
          taskId: null,
          taskSource: null,
          agentType: null,
          status: "spawning",
          spawnedAt: new Date("2026-03-07T09:00:00.000Z"),
          completedAt: null,
          metadata: {},
        },
      ]);

      const result = await service.getTree();

      expect(result).toEqual([
        {
          sessionId: "agent-root",
          parentSessionId: null,
          status: "spawning",
          agentType: null,
          taskSource: null,
          spawnedAt: "2026-03-07T09:00:00.000Z",
          completedAt: null,
        },
      ]);
    });
  });
});