stack/apps/api/src/federation/federation-agent.service.spec.ts

/**
 * Tests for Federation Agent Service
 */

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { HttpService } from "@nestjs/axios";
import { ConfigService } from "@nestjs/config";
import { FederationAgentService } from "./federation-agent.service";
import { CommandService } from "./command.service";
import { PrismaService } from "../prisma/prisma.service";
import { FederationConnectionStatus } from "@prisma/client";
import { of, throwError } from "rxjs";
import type {
  SpawnAgentCommandPayload,
  AgentStatusCommandPayload,
  KillAgentCommandPayload,
  SpawnAgentResponseData,
  AgentStatusResponseData,
  KillAgentResponseData,
} from "./types/federation-agent.types";

describe("FederationAgentService", () => {
  let service: FederationAgentService;
  let commandService: ReturnType<typeof vi.mocked<CommandService>>;
  let prisma: ReturnType<typeof vi.mocked<PrismaService>>;
  let httpService: ReturnType<typeof vi.mocked<HttpService>>;
  let configService: ReturnType<typeof vi.mocked<ConfigService>>;

  const mockWorkspaceId = "workspace-1";
  const mockConnectionId = "connection-1";
  const mockAgentId = "agent-123";
  const mockTaskId = "task-456";
  const mockOrchestratorUrl = "http://localhost:3001";

  beforeEach(async () => {
    const mockCommandService = {
      sendCommand: vi.fn(),
    };

    const mockPrisma = {
      federationConnection: {
        findUnique: vi.fn(),
        findFirst: vi.fn(),
      },
    };

    const mockHttpService = {
      post: vi.fn(),
      get: vi.fn(),
    };

    const mockConfigService = {
      get: vi.fn((key: string) => {
        if (key === "orchestrator.url") {
          return mockOrchestratorUrl;
        }
        return undefined;
      }),
    };

    const module: TestingModule = await Test.createTestingModule({
      providers: [
        FederationAgentService,
        { provide: CommandService, useValue: mockCommandService },
        { provide: PrismaService, useValue: mockPrisma },
        { provide: HttpService, useValue: mockHttpService },
        { provide: ConfigService, useValue: mockConfigService },
      ],
    }).compile();

    service = module.get<FederationAgentService>(FederationAgentService);
    commandService = module.get(CommandService);
    prisma = module.get(PrismaService);
    httpService = module.get(HttpService);
    configService = module.get(ConfigService);
  });

  it("should be defined", () => {
    expect(service).toBeDefined();
  });

  describe("spawnAgentOnRemote", () => {
    const spawnPayload: SpawnAgentCommandPayload = {
      taskId: mockTaskId,
      agentType: "worker",
      context: {
        repository: "git.example.com/org/repo",
        branch: "main",
        workItems: ["item-1"],
      },
    };

    const mockConnection = {
      id: mockConnectionId,
      workspaceId: mockWorkspaceId,
      remoteInstanceId: "remote-instance-1",
      remoteUrl: "https://remote.example.com",
      status: FederationConnectionStatus.ACTIVE,
    };

    it("should spawn agent on remote instance", async () => {
      prisma.federationConnection.findUnique.mockResolvedValue(mockConnection as never);

      const mockCommandResponse = {
        id: "msg-1",
        workspaceId: mockWorkspaceId,
        connectionId: mockConnectionId,
        messageType: "COMMAND" as never,
        messageId: "msg-uuid",
        commandType: "agent.spawn",
        payload: spawnPayload as never,
        response: {
          agentId: mockAgentId,
          status: "spawning",
          spawnedAt: "2026-02-03T14:30:00Z",
        } as never,
        status: "DELIVERED" as never,
        createdAt: new Date(),
        updatedAt: new Date(),
      };

      commandService.sendCommand.mockResolvedValue(mockCommandResponse as never);

      const result = await service.spawnAgentOnRemote(
        mockWorkspaceId,
        mockConnectionId,
        spawnPayload
      );

      expect(prisma.federationConnection.findUnique).toHaveBeenCalledWith({
        where: { id: mockConnectionId, workspaceId: mockWorkspaceId },
      });

      expect(commandService.sendCommand).toHaveBeenCalledWith(
        mockWorkspaceId,
        mockConnectionId,
        "agent.spawn",
        spawnPayload
      );

      expect(result).toEqual(mockCommandResponse);
    });

    it("should throw error if connection not found", async () => {
      prisma.federationConnection.findUnique.mockResolvedValue(null);

      await expect(
        service.spawnAgentOnRemote(mockWorkspaceId, mockConnectionId, spawnPayload)
      ).rejects.toThrow("Connection not found");

      expect(commandService.sendCommand).not.toHaveBeenCalled();
    });

    it("should throw error if connection not active", async () => {
      const inactiveConnection = {
        ...mockConnection,
        status: FederationConnectionStatus.DISCONNECTED,
      };

      prisma.federationConnection.findUnique.mockResolvedValue(inactiveConnection as never);

      await expect(
        service.spawnAgentOnRemote(mockWorkspaceId, mockConnectionId, spawnPayload)
      ).rejects.toThrow("Connection is not active");

      expect(commandService.sendCommand).not.toHaveBeenCalled();
    });
  });

  describe("getAgentStatus", () => {
    const statusPayload: AgentStatusCommandPayload = {
      agentId: mockAgentId,
    };

    const mockConnection = {
      id: mockConnectionId,
      workspaceId: mockWorkspaceId,
      remoteInstanceId: "remote-instance-1",
      remoteUrl: "https://remote.example.com",
      status: FederationConnectionStatus.ACTIVE,
    };

    it("should get agent status from remote instance", async () => {
      prisma.federationConnection.findUnique.mockResolvedValue(mockConnection as never);

      const mockCommandResponse = {
        id: "msg-2",
        workspaceId: mockWorkspaceId,
        connectionId: mockConnectionId,
        messageType: "COMMAND" as never,
        messageId: "msg-uuid-2",
        commandType: "agent.status",
        payload: statusPayload as never,
        response: {
          agentId: mockAgentId,
          taskId: mockTaskId,
          status: "running",
          spawnedAt: "2026-02-03T14:30:00Z",
          startedAt: "2026-02-03T14:30:05Z",
        } as never,
        status: "DELIVERED" as never,
        createdAt: new Date(),
        updatedAt: new Date(),
      };

      commandService.sendCommand.mockResolvedValue(mockCommandResponse as never);

      const result = await service.getAgentStatus(mockWorkspaceId, mockConnectionId, mockAgentId);

      expect(commandService.sendCommand).toHaveBeenCalledWith(
        mockWorkspaceId,
        mockConnectionId,
        "agent.status",
        statusPayload
      );

      expect(result).toEqual(mockCommandResponse);
    });
  });

  describe("killAgentOnRemote", () => {
    const killPayload: KillAgentCommandPayload = {
      agentId: mockAgentId,
    };

    const mockConnection = {
      id: mockConnectionId,
      workspaceId: mockWorkspaceId,
      remoteInstanceId: "remote-instance-1",
      remoteUrl: "https://remote.example.com",
      status: FederationConnectionStatus.ACTIVE,
    };

    it("should kill agent on remote instance", async () => {
      prisma.federationConnection.findUnique.mockResolvedValue(mockConnection as never);

      const mockCommandResponse = {
        id: "msg-3",
        workspaceId: mockWorkspaceId,
        connectionId: mockConnectionId,
        messageType: "COMMAND" as never,
        messageId: "msg-uuid-3",
        commandType: "agent.kill",
        payload: killPayload as never,
        response: {
          agentId: mockAgentId,
          status: "killed",
          killedAt: "2026-02-03T14:35:00Z",
        } as never,
        status: "DELIVERED" as never,
        createdAt: new Date(),
        updatedAt: new Date(),
      };

      commandService.sendCommand.mockResolvedValue(mockCommandResponse as never);

      const result = await service.killAgentOnRemote(
        mockWorkspaceId,
        mockConnectionId,
        mockAgentId
      );

      expect(commandService.sendCommand).toHaveBeenCalledWith(
        mockWorkspaceId,
        mockConnectionId,
        "agent.kill",
        killPayload
      );

      expect(result).toEqual(mockCommandResponse);
    });
  });

  describe("handleAgentCommand", () => {
    const mockConnection = {
      id: mockConnectionId,
      workspaceId: mockWorkspaceId,
      remoteInstanceId: "remote-instance-1",
      remoteUrl: "https://remote.example.com",
      status: FederationConnectionStatus.ACTIVE,
    };

    it("should handle agent.spawn command", async () => {
      const spawnPayload: SpawnAgentCommandPayload = {
        taskId: mockTaskId,
        agentType: "worker",
        context: {
          repository: "git.example.com/org/repo",
          branch: "main",
          workItems: ["item-1"],
        },
      };

      prisma.federationConnection.findFirst.mockResolvedValue(mockConnection as never);

      const mockOrchestratorResponse = {
        agentId: mockAgentId,
        status: "spawning",
      };

      httpService.post.mockReturnValue(
        of({
          data: mockOrchestratorResponse,
          status: 200,
          statusText: "OK",
          headers: {},
          config: {} as never,
        }) as never
      );

      const result = await service.handleAgentCommand(
        "remote-instance-1",
        "agent.spawn",
        spawnPayload
      );

      expect(httpService.post).toHaveBeenCalledWith(
        `${mockOrchestratorUrl}/agents/spawn`,
        expect.objectContaining({
          taskId: mockTaskId,
          agentType: "worker",
        })
      );

      expect(result.success).toBe(true);
      expect(result.data).toEqual({
        agentId: mockAgentId,
        status: "spawning",
        spawnedAt: expect.any(String),
      });
    });

    it("should handle agent.status command", async () => {
      const statusPayload: AgentStatusCommandPayload = {
        agentId: mockAgentId,
      };

      prisma.federationConnection.findFirst.mockResolvedValue(mockConnection as never);

      const mockOrchestratorResponse = {
        agentId: mockAgentId,
        taskId: mockTaskId,
        status: "running",
        spawnedAt: "2026-02-03T14:30:00Z",
        startedAt: "2026-02-03T14:30:05Z",
      };

      httpService.get.mockReturnValue(
        of({
          data: mockOrchestratorResponse,
          status: 200,
          statusText: "OK",
          headers: {},
          config: {} as never,
        }) as never
      );

      const result = await service.handleAgentCommand(
        "remote-instance-1",
        "agent.status",
        statusPayload
      );

      expect(httpService.get).toHaveBeenCalledWith(
        `${mockOrchestratorUrl}/agents/${mockAgentId}/status`
      );

      expect(result.success).toBe(true);
      expect(result.data).toEqual(mockOrchestratorResponse);
    });

    it("should handle agent.kill command", async () => {
      const killPayload: KillAgentCommandPayload = {
        agentId: mockAgentId,
      };

      prisma.federationConnection.findFirst.mockResolvedValue(mockConnection as never);

      const mockOrchestratorResponse = {
        message: `Agent ${mockAgentId} killed successfully`,
      };

      httpService.post.mockReturnValue(
        of({
          data: mockOrchestratorResponse,
          status: 200,
          statusText: "OK",
          headers: {},
          config: {} as never,
        }) as never
      );

      const result = await service.handleAgentCommand(
        "remote-instance-1",
        "agent.kill",
        killPayload
      );

      expect(httpService.post).toHaveBeenCalledWith(
        `${mockOrchestratorUrl}/agents/${mockAgentId}/kill`,
        {}
      );

      expect(result.success).toBe(true);
      expect(result.data).toEqual({
        agentId: mockAgentId,
        status: "killed",
        killedAt: expect.any(String),
      });
    });

    it("should return error for unknown command type", async () => {
      prisma.federationConnection.findFirst.mockResolvedValue(mockConnection as never);

      const result = await service.handleAgentCommand("remote-instance-1", "agent.unknown", {});

      expect(result.success).toBe(false);
      expect(result.error).toContain("Unknown agent command type: agent.unknown");
    });

    it("should throw error if connection not found", async () => {
      prisma.federationConnection.findFirst.mockResolvedValue(null);

      await expect(
        service.handleAgentCommand("remote-instance-1", "agent.spawn", {})
      ).rejects.toThrow("No connection found for remote instance");
    });

    it("should handle orchestrator errors", async () => {
      const spawnPayload: SpawnAgentCommandPayload = {
        taskId: mockTaskId,
        agentType: "worker",
        context: {
          repository: "git.example.com/org/repo",
          branch: "main",
          workItems: ["item-1"],
        },
      };

      prisma.federationConnection.findFirst.mockResolvedValue(mockConnection as never);

      httpService.post.mockReturnValue(
        throwError(() => new Error("Orchestrator connection failed")) as never
      );

      const result = await service.handleAgentCommand(
        "remote-instance-1",
        "agent.spawn",
        spawnPayload
      );

      expect(result.success).toBe(false);
      expect(result.error).toContain("Orchestrator connection failed");
    });
  });
});