stack/apps/orchestrator/src/api/agents/internal-agent.provider.spec.ts

import { beforeEach, describe, expect, it, vi } from "vitest";
import type { AgentConversationMessage, AgentSessionTree } from "@prisma/client";
import { AgentControlService } from "./agent-control.service";
import { AgentMessagesService } from "./agent-messages.service";
import { AgentTreeService } from "./agent-tree.service";
import { InternalAgentProvider } from "./internal-agent.provider";

describe("InternalAgentProvider", () => {
  let provider: InternalAgentProvider;
  let messagesService: {
    getMessages: ReturnType<typeof vi.fn>;
    getReplayMessages: ReturnType<typeof vi.fn>;
    getMessagesAfter: ReturnType<typeof vi.fn>;
  };
  let controlService: {
    injectMessage: ReturnType<typeof vi.fn>;
    pauseAgent: ReturnType<typeof vi.fn>;
    resumeAgent: ReturnType<typeof vi.fn>;
    killAgent: ReturnType<typeof vi.fn>;
  };
  let treeService: {
    listSessions: ReturnType<typeof vi.fn>;
    getSession: ReturnType<typeof vi.fn>;
  };

  beforeEach(() => {
    messagesService = {
      getMessages: vi.fn(),
      getReplayMessages: vi.fn(),
      getMessagesAfter: vi.fn(),
    };

    controlService = {
      injectMessage: vi.fn().mockResolvedValue(undefined),
      pauseAgent: vi.fn().mockResolvedValue(undefined),
      resumeAgent: vi.fn().mockResolvedValue(undefined),
      killAgent: vi.fn().mockResolvedValue(undefined),
    };

    treeService = {
      listSessions: vi.fn(),
      getSession: vi.fn(),
    };

    provider = new InternalAgentProvider(
      messagesService as unknown as AgentMessagesService,
      controlService as unknown as AgentControlService,
      treeService as unknown as AgentTreeService
    );
  });

  it("maps paginated sessions", async () => {
    const sessionEntry: AgentSessionTree = {
      id: "tree-1",
      sessionId: "session-1",
      parentSessionId: "parent-1",
      provider: "internal",
      missionId: null,
      taskId: "task-123",
      taskSource: "queue",
      agentType: "worker",
      status: "running",
      spawnedAt: new Date("2026-03-07T10:00:00.000Z"),
      completedAt: null,
      metadata: { branch: "feat/test" },
    };

    treeService.listSessions.mockResolvedValue({
      sessions: [sessionEntry],
      total: 1,
      cursor: "next-cursor",
    });

    const result = await provider.listSessions("cursor-1", 25);

    expect(treeService.listSessions).toHaveBeenCalledWith("cursor-1", 25);
    expect(result).toEqual({
      sessions: [
        {
          id: "session-1",
          providerId: "internal",
          providerType: "internal",
          label: "task-123",
          status: "active",
          parentSessionId: "parent-1",
          createdAt: new Date("2026-03-07T10:00:00.000Z"),
          updatedAt: new Date("2026-03-07T10:00:00.000Z"),
          metadata: { branch: "feat/test" },
        },
      ],
      total: 1,
      cursor: "next-cursor",
    });
  });

  it("returns null for missing session", async () => {
    treeService.getSession.mockResolvedValue(null);

    const result = await provider.getSession("missing-session");

    expect(treeService.getSession).toHaveBeenCalledWith("missing-session");
    expect(result).toBeNull();
  });

  it("maps message history and parses skip cursor", async () => {
    const message: AgentConversationMessage = {
      id: "msg-1",
      sessionId: "session-1",
      provider: "internal",
      role: "agent",
      content: "hello",
      timestamp: new Date("2026-03-07T10:05:00.000Z"),
      metadata: { tokens: 42 },
    };

    messagesService.getMessages.mockResolvedValue({
      messages: [message],
      total: 10,
    });

    const result = await provider.getMessages("session-1", 30, "2");

    expect(messagesService.getMessages).toHaveBeenCalledWith("session-1", 30, 2);
    expect(result).toEqual([
      {
        id: "msg-1",
        sessionId: "session-1",
        role: "assistant",
        content: "hello",
        timestamp: new Date("2026-03-07T10:05:00.000Z"),
        metadata: { tokens: 42 },
      },
    ]);
  });

  it("routes control operations through AgentControlService", async () => {
    const injectResult = await provider.injectMessage("session-1", "new instruction");

    await provider.pauseSession("session-1");
    await provider.resumeSession("session-1");
    await provider.killSession("session-1", false);

    expect(controlService.injectMessage).toHaveBeenCalledWith(
      "session-1",
      "internal-provider",
      "new instruction"
    );
    expect(injectResult).toEqual({ accepted: true });
    expect(controlService.pauseAgent).toHaveBeenCalledWith("session-1", "internal-provider");
    expect(controlService.resumeAgent).toHaveBeenCalledWith("session-1", "internal-provider");
    expect(controlService.killAgent).toHaveBeenCalledWith("session-1", "internal-provider", false);
  });

  it("streams replay and incremental messages", async () => {
    const replayMessage: AgentConversationMessage = {
      id: "msg-replay",
      sessionId: "session-1",
      provider: "internal",
      role: "agent",
      content: "replay",
      timestamp: new Date("2026-03-07T10:00:00.000Z"),
      metadata: {},
    };
    const incrementalMessage: AgentConversationMessage = {
      id: "msg-live",
      sessionId: "session-1",
      provider: "internal",
      role: "operator",
      content: "live",
      timestamp: new Date("2026-03-07T10:00:01.000Z"),
      metadata: {},
    };

    messagesService.getReplayMessages.mockResolvedValue([replayMessage]);
    messagesService.getMessagesAfter
      .mockResolvedValueOnce([incrementalMessage])
      .mockResolvedValueOnce([]);

    const iterator = provider.streamMessages("session-1")[Symbol.asyncIterator]();

    const first = await iterator.next();
    const second = await iterator.next();

    expect(first.done).toBe(false);
    expect(first.value).toEqual({
      id: "msg-replay",
      sessionId: "session-1",
      role: "assistant",
      content: "replay",
      timestamp: new Date("2026-03-07T10:00:00.000Z"),
      metadata: {},
    });
    expect(second.done).toBe(false);
    expect(second.value).toEqual({
      id: "msg-live",
      sessionId: "session-1",
      role: "user",
      content: "live",
      timestamp: new Date("2026-03-07T10:00:01.000Z"),
      metadata: {},
    });

    await iterator.return?.();

    expect(messagesService.getReplayMessages).toHaveBeenCalledWith("session-1", 50);
    expect(messagesService.getMessagesAfter).toHaveBeenCalledWith(
      "session-1",
      new Date("2026-03-07T10:00:00.000Z"),
      "msg-replay"
    );
  });

  it("reports provider availability", async () => {
    await expect(provider.isAvailable()).resolves.toBe(true);
  });
});