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

import { Logger } from "@nestjs/common";
import type {
  AgentMessage,
  AgentSession,
  AgentSessionList,
  IAgentProvider,
  InjectResult,
} from "@mosaic/shared";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { AgentProviderRegistry } from "./agent-provider.registry";
import { InternalAgentProvider } from "./internal-agent.provider";

type MockProvider = IAgentProvider & {
  listSessions: ReturnType<typeof vi.fn>;
  getSession: ReturnType<typeof vi.fn>;
};

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

const createProvider = (providerId: string, sessions: AgentSession[] = []): MockProvider => {
  return {
    providerId,
    providerType: providerId,
    displayName: providerId,
    listSessions: vi.fn().mockResolvedValue({
      sessions,
      total: sessions.length,
    } as AgentSessionList),
    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("AgentProviderRegistry", () => {
  let registry: AgentProviderRegistry;
  let internalProvider: MockProvider;

  beforeEach(() => {
    internalProvider = createProvider("internal");
    registry = new AgentProviderRegistry(internalProvider as unknown as InternalAgentProvider);
  });

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

  it("registers InternalAgentProvider on module init", () => {
    registry.onModuleInit();

    expect(registry.getProvider("internal")).toBe(internalProvider);
  });

  it("registers providers and returns null for unknown provider ids", () => {
    const externalProvider = createProvider("openclaw");

    registry.registerProvider(externalProvider);

    expect(registry.getProvider("openclaw")).toBe(externalProvider);
    expect(registry.getProvider("missing")).toBeNull();
  });

  it("aggregates and sorts sessions from all providers", async () => {
    const internalSessions: AgentSession[] = [
      {
        id: "session-older",
        providerId: "internal",
        providerType: "internal",
        status: "active",
        createdAt: new Date("2026-03-07T10:00:00.000Z"),
        updatedAt: new Date("2026-03-07T10:10:00.000Z"),
      },
    ];

    const externalSessions: AgentSession[] = [
      {
        id: "session-newer",
        providerId: "openclaw",
        providerType: "external",
        status: "paused",
        createdAt: new Date("2026-03-07T09:00:00.000Z"),
        updatedAt: new Date("2026-03-07T10:20:00.000Z"),
      },
    ];

    internalProvider.listSessions.mockResolvedValue({
      sessions: internalSessions,
      total: internalSessions.length,
    } as AgentSessionList);

    const externalProvider = createProvider("openclaw", externalSessions);
    registry.onModuleInit();
    registry.registerProvider(externalProvider);

    const result = await registry.listAllSessions();

    expect(result.map((session) => session.id)).toEqual(["session-newer", "session-older"]);
    expect(internalProvider.listSessions).toHaveBeenCalledTimes(1);
    expect(externalProvider.listSessions).toHaveBeenCalledTimes(1);
  });

  it("skips provider failures and logs warning", async () => {
    const warnSpy = vi.spyOn(Logger.prototype, "warn").mockImplementation(() => undefined);

    const healthyProvider = createProvider("healthy", [
      {
        id: "session-1",
        providerId: "healthy",
        providerType: "external",
        status: "active",
        createdAt: new Date("2026-03-07T11:00:00.000Z"),
        updatedAt: new Date("2026-03-07T11:00:00.000Z"),
      },
    ]);

    const failingProvider = createProvider("failing");
    failingProvider.listSessions.mockRejectedValue(new Error("provider offline"));

    registry.onModuleInit();
    registry.registerProvider(healthyProvider);
    registry.registerProvider(failingProvider);

    const result = await registry.listAllSessions();

    expect(result).toHaveLength(1);
    expect(result[0]?.id).toBe("session-1");
    expect(warnSpy).toHaveBeenCalledWith(
      expect.stringContaining("Failed to list sessions for provider failing")
    );
  });

  it("finds a provider for an existing session", async () => {
    const targetSession: AgentSession = {
      id: "session-found",
      providerId: "openclaw",
      providerType: "external",
      status: "active",
      createdAt: new Date("2026-03-07T12:00:00.000Z"),
      updatedAt: new Date("2026-03-07T12:10:00.000Z"),
    };

    const openclawProvider = createProvider("openclaw");
    openclawProvider.getSession.mockResolvedValue(targetSession);

    registry.onModuleInit();
    registry.registerProvider(openclawProvider);

    const result = await registry.getProviderForSession(targetSession.id);

    expect(result).toEqual({
      provider: openclawProvider,
      session: targetSession,
    });
    expect(internalProvider.getSession).toHaveBeenCalledWith(targetSession.id);
    expect(openclawProvider.getSession).toHaveBeenCalledWith(targetSession.id);
  });

  it("returns null when no provider has the requested session", async () => {
    const openclawProvider = createProvider("openclaw");

    registry.onModuleInit();
    registry.registerProvider(openclawProvider);

    await expect(registry.getProviderForSession("missing-session")).resolves.toBeNull();
  });

  it("continues searching providers when getSession throws", async () => {
    const warnSpy = vi.spyOn(Logger.prototype, "warn").mockImplementation(() => undefined);
    const failingProvider = createProvider("failing");
    failingProvider.getSession.mockRejectedValue(new Error("provider timeout"));

    const healthySession: AgentSession = {
      id: "session-healthy",
      providerId: "healthy",
      providerType: "external",
      status: "active",
      createdAt: new Date("2026-03-07T12:15:00.000Z"),
      updatedAt: new Date("2026-03-07T12:16:00.000Z"),
    };

    const healthyProvider = createProvider("healthy");
    healthyProvider.getSession.mockResolvedValue(healthySession);

    registry.onModuleInit();
    registry.registerProvider(failingProvider);
    registry.registerProvider(healthyProvider);

    const result = await registry.getProviderForSession(healthySession.id);

    expect(result).toEqual({ provider: healthyProvider, session: healthySession });
    expect(warnSpy).toHaveBeenCalledWith(
      expect.stringContaining("Failed to get session session-healthy for provider failing")
    );
  });
});