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

import { NotFoundException } from "@nestjs/common";
import { beforeEach, describe, expect, it, vi } from "vitest";
import type { AgentMessage, AgentSession, IAgentProvider, InjectResult } from "@mosaic/shared";
import type { PrismaService } from "../../prisma/prisma.service";
import { AgentProviderRegistry } from "../agents/agent-provider.registry";
import { MissionControlService } from "./mission-control.service";

type MockProvider = IAgentProvider & {
  listSessions: ReturnType<typeof vi.fn>;
  getSession: ReturnType<typeof vi.fn>;
  getMessages: ReturnType<typeof vi.fn>;
  injectMessage: ReturnType<typeof vi.fn>;
  pauseSession: ReturnType<typeof vi.fn>;
  resumeSession: ReturnType<typeof vi.fn>;
  killSession: ReturnType<typeof vi.fn>;
  streamMessages: ReturnType<typeof vi.fn>;
};

const emptyMessageStream = async function* (): AsyncIterable<AgentMessage> {
  return;
};

const createProvider = (providerId = "internal"): MockProvider => ({
  providerId,
  providerType: providerId,
  displayName: providerId,
  listSessions: vi.fn().mockResolvedValue({ sessions: [], total: 0 }),
  getSession: vi.fn().mockResolvedValue(null),
  getMessages: vi.fn().mockResolvedValue([]),
  injectMessage: vi.fn().mockResolvedValue({ accepted: true } as InjectResult),
  pauseSession: vi.fn().mockResolvedValue(undefined),
  resumeSession: vi.fn().mockResolvedValue(undefined),
  killSession: vi.fn().mockResolvedValue(undefined),
  streamMessages: vi.fn().mockReturnValue(emptyMessageStream()),
  isAvailable: vi.fn().mockResolvedValue(true),
});

describe("MissionControlService", () => {
  let service: MissionControlService;
  let registry: {
    listAllSessions: ReturnType<typeof vi.fn>;
    getProviderForSession: ReturnType<typeof vi.fn>;
  };
  let prisma: {
    operatorAuditLog: {
      create: ReturnType<typeof vi.fn>;
    };
  };

  const session: AgentSession = {
    id: "session-1",
    providerId: "internal",
    providerType: "internal",
    status: "active",
    createdAt: new Date("2026-03-07T14:00:00.000Z"),
    updatedAt: new Date("2026-03-07T14:01:00.000Z"),
  };

  beforeEach(() => {
    registry = {
      listAllSessions: vi.fn().mockResolvedValue([session]),
      getProviderForSession: vi.fn().mockResolvedValue(null),
    };

    prisma = {
      operatorAuditLog: {
        create: vi.fn().mockResolvedValue(undefined),
      },
    };

    service = new MissionControlService(
      registry as unknown as AgentProviderRegistry,
      prisma as unknown as PrismaService
    );
  });

  it("lists sessions from the registry", async () => {
    await expect(service.listSessions()).resolves.toEqual([session]);
    expect(registry.listAllSessions).toHaveBeenCalledTimes(1);
  });

  it("returns a session when it is found", async () => {
    const provider = createProvider("internal");
    registry.getProviderForSession.mockResolvedValue({ provider, session });

    await expect(service.getSession(session.id)).resolves.toEqual(session);
  });

  it("throws NotFoundException when session lookup fails", async () => {
    await expect(service.getSession("missing-session")).rejects.toBeInstanceOf(NotFoundException);
  });

  it("gets messages from the resolved provider", async () => {
    const provider = createProvider("openclaw");
    const messages: AgentMessage[] = [
      {
        id: "message-1",
        sessionId: session.id,
        role: "assistant",
        content: "hello",
        timestamp: new Date("2026-03-07T14:01:00.000Z"),
      },
    ];

    provider.getMessages.mockResolvedValue(messages);
    registry.getProviderForSession.mockResolvedValue({ provider, session });

    await expect(service.getMessages(session.id, 25, "10")).resolves.toEqual(messages);
    expect(provider.getMessages).toHaveBeenCalledWith(session.id, 25, "10");
  });

  it("injects a message and writes an audit log", async () => {
    const provider = createProvider("internal");
    const injectResult: InjectResult = { accepted: true, messageId: "msg-1" };
    provider.injectMessage.mockResolvedValue(injectResult);
    registry.getProviderForSession.mockResolvedValue({ provider, session });

    await expect(service.injectMessage(session.id, "ship it", "operator-1")).resolves.toEqual(
      injectResult
    );

    expect(provider.injectMessage).toHaveBeenCalledWith(session.id, "ship it");
    expect(prisma.operatorAuditLog.create).toHaveBeenCalledWith({
      data: {
        sessionId: session.id,
        userId: "operator-1",
        provider: "internal",
        action: "inject",
        content: "ship it",
        metadata: {
          payload: { message: "ship it" },
        },
      },
    });
  });

  it("pauses and resumes using default operator id", async () => {
    const provider = createProvider("openclaw");
    registry.getProviderForSession.mockResolvedValue({ provider, session });

    await service.pauseSession(session.id);
    await service.resumeSession(session.id);

    expect(provider.pauseSession).toHaveBeenCalledWith(session.id);
    expect(provider.resumeSession).toHaveBeenCalledWith(session.id);
    expect(prisma.operatorAuditLog.create).toHaveBeenNthCalledWith(1, {
      data: {
        sessionId: session.id,
        userId: "mission-control",
        provider: "openclaw",
        action: "pause",
        metadata: {
          payload: {},
        },
      },
    });
    expect(prisma.operatorAuditLog.create).toHaveBeenNthCalledWith(2, {
      data: {
        sessionId: session.id,
        userId: "mission-control",
        provider: "openclaw",
        action: "resume",
        metadata: {
          payload: {},
        },
      },
    });
  });

  it("kills with provided force value and writes audit log", async () => {
    const provider = createProvider("openclaw");
    registry.getProviderForSession.mockResolvedValue({ provider, session });

    await service.killSession(session.id, false, "operator-2");

    expect(provider.killSession).toHaveBeenCalledWith(session.id, false);
    expect(prisma.operatorAuditLog.create).toHaveBeenCalledWith({
      data: {
        sessionId: session.id,
        userId: "operator-2",
        provider: "openclaw",
        action: "kill",
        metadata: {
          payload: { force: false },
        },
      },
    });
  });

  it("resolves provider message stream", async () => {
    const provider = createProvider("internal");
    const messageStream = (async function* (): AsyncIterable<AgentMessage> {
      yield {
        id: "message-1",
        sessionId: session.id,
        role: "assistant",
        content: "stream",
        timestamp: new Date("2026-03-07T14:03:00.000Z"),
      };
    })();

    provider.streamMessages.mockReturnValue(messageStream);
    registry.getProviderForSession.mockResolvedValue({ provider, session });

    await expect(service.streamMessages(session.id)).resolves.toBe(messageStream);
    expect(provider.streamMessages).toHaveBeenCalledWith(session.id);
  });

  it("does not write audit log when session cannot be resolved", async () => {
    await expect(service.pauseSession("missing-session")).rejects.toBeInstanceOf(NotFoundException);
    expect(prisma.operatorAuditLog.create).not.toHaveBeenCalled();
  });
});