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

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { AgentControlService } from "./agent-control.service";
import { PrismaService } from "../../prisma/prisma.service";
import { KillswitchService } from "../../killswitch/killswitch.service";

describe("AgentControlService", () => {
  let service: AgentControlService;
  let prisma: {
    agentSessionTree: {
      findUnique: ReturnType<typeof vi.fn>;
      updateMany: ReturnType<typeof vi.fn>;
    };
    agentConversationMessage: {
      create: ReturnType<typeof vi.fn>;
    };
    operatorAuditLog: {
      create: ReturnType<typeof vi.fn>;
    };
  };
  let killswitchService: {
    killAgent: ReturnType<typeof vi.fn>;
  };

  beforeEach(() => {
    prisma = {
      agentSessionTree: {
        findUnique: vi.fn(),
        updateMany: vi.fn().mockResolvedValue({ count: 1 }),
      },
      agentConversationMessage: {
        create: vi.fn().mockResolvedValue(undefined),
      },
      operatorAuditLog: {
        create: vi.fn().mockResolvedValue(undefined),
      },
    };

    killswitchService = {
      killAgent: vi.fn().mockResolvedValue(undefined),
    };

    service = new AgentControlService(
      prisma as unknown as PrismaService,
      killswitchService as unknown as KillswitchService
    );
  });

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

  describe("injectMessage", () => {
    it("creates conversation message and audit log when tree entry exists", async () => {
      prisma.agentSessionTree.findUnique.mockResolvedValue({ id: "tree-1" });

      await service.injectMessage("agent-123", "operator-abc", "Please continue");

      expect(prisma.agentSessionTree.findUnique).toHaveBeenCalledWith({
        where: { sessionId: "agent-123" },
        select: { id: true },
      });
      expect(prisma.agentConversationMessage.create).toHaveBeenCalledWith({
        data: {
          sessionId: "agent-123",
          role: "operator",
          content: "Please continue",
          provider: "internal",
          metadata: {},
        },
      });
      expect(prisma.operatorAuditLog.create).toHaveBeenCalledWith({
        data: {
          sessionId: "agent-123",
          userId: "operator-abc",
          provider: "internal",
          action: "inject",
          metadata: {
            payload: {
              message: "Please continue",
            },
          },
        },
      });
    });

    it("creates only audit log when no tree entry exists", async () => {
      prisma.agentSessionTree.findUnique.mockResolvedValue(null);

      await service.injectMessage("agent-456", "operator-def", "Nudge message");

      expect(prisma.agentConversationMessage.create).not.toHaveBeenCalled();
      expect(prisma.operatorAuditLog.create).toHaveBeenCalledWith({
        data: {
          sessionId: "agent-456",
          userId: "operator-def",
          provider: "internal",
          action: "inject",
          metadata: {
            payload: {
              message: "Nudge message",
            },
          },
        },
      });
    });
  });

  describe("pauseAgent", () => {
    it("updates tree status to paused and creates audit log", async () => {
      await service.pauseAgent("agent-789", "operator-pause");

      expect(prisma.agentSessionTree.updateMany).toHaveBeenCalledWith({
        where: { sessionId: "agent-789" },
        data: { status: "paused" },
      });
      expect(prisma.operatorAuditLog.create).toHaveBeenCalledWith({
        data: {
          sessionId: "agent-789",
          userId: "operator-pause",
          provider: "internal",
          action: "pause",
          metadata: {
            payload: {},
          },
        },
      });
    });
  });

  describe("resumeAgent", () => {
    it("updates tree status to running and creates audit log", async () => {
      await service.resumeAgent("agent-321", "operator-resume");

      expect(prisma.agentSessionTree.updateMany).toHaveBeenCalledWith({
        where: { sessionId: "agent-321" },
        data: { status: "running" },
      });
      expect(prisma.operatorAuditLog.create).toHaveBeenCalledWith({
        data: {
          sessionId: "agent-321",
          userId: "operator-resume",
          provider: "internal",
          action: "resume",
          metadata: {
            payload: {},
          },
        },
      });
    });
  });

  describe("killAgent", () => {
    it("delegates kill to killswitch and logs audit", async () => {
      await service.killAgent("agent-654", "operator-kill", false);

      expect(killswitchService.killAgent).toHaveBeenCalledWith("agent-654");
      expect(prisma.operatorAuditLog.create).toHaveBeenCalledWith({
        data: {
          sessionId: "agent-654",
          userId: "operator-kill",
          provider: "internal",
          action: "kill",
          metadata: {
            payload: {
              force: false,
            },
          },
        },
      });
    });
  });
});