fusion/packages/dashboard/src/__tests__/websocket.test.ts

import { EventEmitter, once } from "node:events";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { WebSocket } from "ws";
import type { Task } from "@fusion/core";
import { createServer } from "../server.js";
import { WebSocketManager } from "../websocket.js";
import { InMemoryBadgePubSub, type BadgePubSub } from "../badge-pubsub.js";
import { createLoopbackIntegrationTest } from "./loopback-integration-test.js";

const websocketIntegrationTest = await createLoopbackIntegrationTest("websocket /api/ws integration");

class MockSocket extends EventEmitter {
  readyState: number = WebSocket.OPEN;
  sent: string[] = [];
  ping = vi.fn();
  terminate = vi.fn(() => {
    this.readyState = WebSocket.CLOSED;
  });
  send = vi.fn((payload: string) => {
    this.sent.push(payload);
  });
  close = vi.fn(() => {
    this.readyState = WebSocket.CLOSED;
    this.emit("close");
  });
}

class MockStore extends EventEmitter {
  task: Task;

  constructor(task: Task) {
    super();
    this.task = task;
  }

  getRootDir(): string {
    return process.cwd();
  }

  getFusionDir(): string {
    return process.cwd() + "/.fusion";
  }

  getDatabase() {
    return {
      exec: vi.fn(),
      prepare: vi.fn().mockReturnValue({ run: vi.fn().mockReturnValue({ changes: 0 }), get: vi.fn(), all: vi.fn().mockReturnValue([]) }),
    };
  }

  async listTasks(): Promise<Task[]> {
    return [this.task];
  }

  async getTask(id: string): Promise<Task> {
    if (id !== this.task.id) {
      const error = Object.assign(new Error("Task not found"), { code: "ENOENT" });
      throw error;
    }

    return this.task;
  }

  getMissionStore() {
    // Return a mock mission store that has minimal functionality for the tests
    return {
      listMissions: vi.fn().mockResolvedValue([]),
      getMission: vi.fn(),
      createMission: vi.fn(),
      updateMission: vi.fn(),
      deleteMission: vi.fn(),
      getMissionWithHierarchy: vi.fn(),
      listMilestones: vi.fn().mockResolvedValue([]),
      getMilestone: vi.fn(),
      addMilestone: vi.fn(),
      updateMilestone: vi.fn(),
      deleteMilestone: vi.fn(),
      reorderMilestones: vi.fn(),
      listSlices: vi.fn().mockResolvedValue([]),
      getSlice: vi.fn(),
      addSlice: vi.fn(),
      updateSlice: vi.fn(),
      deleteSlice: vi.fn(),
      reorderSlices: vi.fn(),
      activateSlice: vi.fn(),
      listFeatures: vi.fn().mockResolvedValue([]),
      getFeature: vi.fn(),
      addFeature: vi.fn(),
      updateFeature: vi.fn(),
      deleteFeature: vi.fn(),
      linkFeatureToTask: vi.fn(),
      unlinkFeatureFromTask: vi.fn(),
      getFeatureRollups: vi.fn().mockResolvedValue([]),
    };
  }
}

function createTask(overrides: Partial<Task> = {}): Task {
  return {
    id: "FN-063",
    title: "Realtime badge updates",
    description: "Test task",
    column: "in-review",
    dependencies: [],
    steps: [],
    currentStep: 0,
    log: [],
    createdAt: "2026-03-30T00:00:00.000Z",
    updatedAt: "2026-03-30T00:00:00.000Z",
    columnMovedAt: "2026-03-30T00:00:00.000Z",
    prInfo: {
      url: "https://github.com/owner/repo/pull/42",
      number: 42,
      status: "open",
      title: "Tracked PR",
      headBranch: "feature/test",
      baseBranch: "main",
      commentCount: 0,
      lastCheckedAt: "2026-03-30T00:00:00.000Z",
    },
    ...overrides,
  };
}

async function waitForExpectation(assertion: () => void, timeoutMs: number = 1_000): Promise<void> {
  const start = Date.now();
  while (true) {
    try {
      assertion();
      return;
    } catch (error) {
      if (Date.now() - start >= timeoutMs) {
        throw error;
      }
      await new Promise((resolve) => setTimeout(resolve, 10));
    }
  }
}

describe("WebSocketManager", () => {
  beforeEach(() => {
    vi.useFakeTimers();
  });

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

  it("routes subscribe and unsubscribe messages to task channels", () => {
    const manager = new WebSocketManager();
    const socket = new MockSocket();

    manager.addClient(socket as unknown as WebSocket, "client-1");
    socket.emit("message", Buffer.from(JSON.stringify({ type: "subscribe", taskId: "FN-063" })));
    expect(manager.getSubscriptionCount("FN-063")).toBe(1);

    socket.emit("message", Buffer.from(JSON.stringify({ type: "unsubscribe", taskId: "FN-063" })));
    expect(manager.getSubscriptionCount("FN-063")).toBe(0);
  });

  it("broadcasts badge updates only to subscribed clients", () => {
    const manager = new WebSocketManager();
    const first = new MockSocket();
    const second = new MockSocket();

    manager.addClient(first as unknown as WebSocket, "client-1");
    manager.addClient(second as unknown as WebSocket, "client-2");

    first.emit("message", Buffer.from(JSON.stringify({ type: "subscribe", taskId: "FN-063" })));
    second.emit("message", Buffer.from(JSON.stringify({ type: "subscribe", taskId: "FN-064" })));

    manager.broadcastBadgeUpdate("FN-063", {
      prInfo: null,
      issueInfo: {
        url: "https://github.com/owner/repo/issues/2",
        number: 2,
        state: "closed",
        title: "Issue",
        stateReason: "completed",
      },
      timestamp: "2026-03-30T12:00:00.000Z",
    });

    expect(first.send).toHaveBeenCalledTimes(1);
    expect(second.send).not.toHaveBeenCalled();
    expect(JSON.parse(first.sent[0])).toMatchObject({
      type: "badge:updated",
      taskId: "FN-063",
      prInfo: null,
      issueInfo: { number: 2 },
    });
  });

  it("keeps connections alive when pong responses arrive", () => {
    const manager = new WebSocketManager({ heartbeatIntervalMs: 100 });
    const socket = new MockSocket();

    manager.addClient(socket as unknown as WebSocket, "client-1");

    vi.advanceTimersByTime(100);
    expect(socket.ping).toHaveBeenCalledTimes(1);

    socket.emit("pong");
    vi.advanceTimersByTime(100);

    expect(socket.terminate).not.toHaveBeenCalled();
    expect(manager.getClientCount()).toBe(1);
  });

  it("terminates dead connections and cleans up subscriptions", () => {
    const manager = new WebSocketManager({ heartbeatIntervalMs: 100 });
    const socket = new MockSocket();

    manager.addClient(socket as unknown as WebSocket, "client-1");
    socket.emit("message", Buffer.from(JSON.stringify({ type: "subscribe", taskId: "FN-063" })));

    vi.advanceTimersByTime(200);

    expect(socket.terminate).toHaveBeenCalled();
    expect(manager.getClientCount()).toBe(0);
    expect(manager.getSubscriptionCount("FN-063")).toBe(0);
  });

  it("disposes sockets without leaking tracked clients", () => {
    const manager = new WebSocketManager();
    const socket = new MockSocket();

    manager.addClient(socket as unknown as WebSocket, "client-1");
    manager.dispose();

    expect(socket.terminate).toHaveBeenCalled();
    expect(manager.getClientCount()).toBe(0);
    expect(manager.getSubscribedTaskIds()).toEqual([]);
  });

  describe("project-scoped channel isolation", () => {
    it("subscribe uses scoped channel keys - same task ID in different projects", () => {
      const manager = new WebSocketManager();
      const socketA = new MockSocket();
      const socketB = new MockSocket();

      // Add clients bound to different project scopes
      manager.addClient(socketA as unknown as WebSocket, "client-1", "project-a");
      manager.addClient(socketB as unknown as WebSocket, "client-2", "project-b");

      // Subscribe to the same task ID in different projects
      manager.subscribe("client-1", "FN-063", "project-a");
      manager.subscribe("client-2", "FN-063", "project-b");

      // Each project should have its own subscription count
      expect(manager.getSubscriptionCount("FN-063", "project-a")).toBe(1);
      expect(manager.getSubscriptionCount("FN-063", "project-b")).toBe(1);

      // Unscoped query should return 0 (no unscoped subscriptions)
      expect(manager.getSubscriptionCount("FN-063")).toBe(0);
    });

    it("broadcastBadgeUpdate does not leak across projects with same task ID", () => {
      const manager = new WebSocketManager();
      const clientA = new MockSocket();
      const clientB = new MockSocket();

      manager.addClient(clientA as unknown as WebSocket, "client-a", "project-a");
      manager.addClient(clientB as unknown as WebSocket, "client-b", "project-b");

      // Subscribe both to the same task ID in their respective projects
      manager.subscribe("client-a", "FN-063", "project-a");
      manager.subscribe("client-b", "FN-063", "project-b");

      // Broadcast for project-a only
      manager.broadcastBadgeUpdate("FN-063", {
        prInfo: { url: "https://github.com/owner/repo/pull/1", number: 1, status: "merged", title: "Project A PR", headBranch: "feat", baseBranch: "main", commentCount: 0 },
        timestamp: "2026-03-30T12:00:00.000Z",
      }, "project-a");

      // Only project-a client should receive the update
      expect(clientA.send).toHaveBeenCalledTimes(1);
      expect(clientB.send).not.toHaveBeenCalled();

      const messageA = JSON.parse(clientA.sent[0]);
      expect(messageA).toMatchObject({
        type: "badge:updated",
        taskId: "FN-063",
        prInfo: { status: "merged", title: "Project A PR" },
      });
    });

    it("unsubscribe is scoped to project", () => {
      const manager = new WebSocketManager();
      const socket = new MockSocket();

      manager.addClient(socket as unknown as WebSocket, "client-1", "project-a");

      // Subscribe in two different project scopes
      manager.subscribe("client-1", "FN-063", "project-a");
      manager.subscribe("client-1", "FN-063", "project-b");

      // project-a subscription should exist
      expect(manager.getSubscriptionCount("FN-063", "project-a")).toBe(1);
      // project-b subscription should also exist
      expect(manager.getSubscriptionCount("FN-063", "project-b")).toBe(1);

      // Unsubscribe from project-a only
      manager.unsubscribe("client-1", "FN-063", "project-a");

      // project-a should be gone, project-b should remain
      expect(manager.getSubscriptionCount("FN-063", "project-a")).toBe(0);
      expect(manager.getSubscriptionCount("FN-063", "project-b")).toBe(1);
    });

    it("getSubscribedTaskIds returns scoped results", () => {
      const manager = new WebSocketManager();
      const socket = new MockSocket();

      manager.addClient(socket as unknown as WebSocket, "client-1", "project-x");

      // Subscribe to different tasks in different projects
      manager.subscribe("client-1", "FN-001", "project-x");
      manager.subscribe("client-1", "FN-002", "project-x");
      manager.subscribe("client-1", "FN-003", "project-y");

      // Get task IDs for project-x only
      const projectXTasks = manager.getSubscribedTaskIds("project-x");
      expect(projectXTasks).toContain("FN-001");
      expect(projectXTasks).toContain("FN-002");
      expect(projectXTasks).not.toContain("FN-003");

      // Get task IDs for project-y only
      const projectYTasks = manager.getSubscribedTaskIds("project-y");
      expect(projectYTasks).toContain("FN-003");
      expect(projectYTasks).not.toContain("FN-001");
      expect(projectYTasks).not.toContain("FN-002");
    });
  });
});

describe("/api/ws integration", () => {
  websocketIntegrationTest("delivers badge updates to subscribed websocket clients via task:updated events", async () => {
    const initialTask = createTask();
    const store = new MockStore(initialTask);
    const app = createServer(store as any, { githubToken: "test-token" });
    const server = app.listen(0);
    await once(server, "listening");

    const port = (server.address() as import("node:net").AddressInfo).port;
    const client = new WebSocket(`ws://127.0.0.1:${port}/api/ws`);
    const messages: any[] = [];

    client.on("message", (payload) => {
      messages.push(JSON.parse(payload.toString()));
    });

    await once(client, "open");
    client.send(JSON.stringify({ type: "subscribe", taskId: initialTask.id }));

    // Wait for subscription to be established
    await new Promise((resolve) => setTimeout(resolve, 100));

    const updatedTask = createTask({
      prInfo: {
        ...initialTask.prInfo!,
        status: "merged",
        title: "Merged PR",
        lastCheckedAt: "2026-03-30T12:00:00.000Z",
      },
      updatedAt: "2026-03-30T12:00:00.000Z",
    });
    (store as unknown as MockStore).task = updatedTask;
    (store as unknown as MockStore).emit("task:updated", updatedTask);

    await waitForExpectation(() => {
      expect(messages).toContainEqual(expect.objectContaining({
        type: "badge:updated",
        taskId: updatedTask.id,
        prInfo: expect.objectContaining({ status: "merged", title: "Merged PR" }),
        issueInfo: null,
      }));
    });

    client.close();
    await once(client, "close");
    server.close();
    await once(server, "close");
  });
});

/**
 * Multi-instance integration tests for cross-instance badge delivery.
 * These tests verify that badge updates flow correctly between multiple
 * dashboard instances using a shared pub/sub adapter.
 */
describe("multi-instance /api/ws integration", () => {
  websocketIntegrationTest("delivers badge updates from instance A to subscribed client on instance B", async () => {
    // Create a shared pub/sub adapter that both instances will use
    const sharedPubSub: BadgePubSub = new InMemoryBadgePubSub();
    await sharedPubSub.start();

    // Create two separate stores (simulating separate instances)
    const taskA = createTask({ 
      id: "FN-MULTI-001", 
      prInfo: { 
        url: "https://github.com/owner/repo/pull/1", 
        number: 1, 
        status: "open", 
        title: "Original PR",
        headBranch: "feature",
        baseBranch: "main",
        commentCount: 0,
        lastCheckedAt: "2026-03-30T00:00:00.000Z",
      },
    });
    
    const storeA = new MockStore(taskA);
    const storeB = new MockStore(taskA); // Instance B starts with same task data

    // Create server A with zero local websocket subscribers (no local ws clients)
    const appA = createServer(storeA as any, { 
      githubToken: "test-token",
      badgePubSub: sharedPubSub,
    });
    const serverA = appA.listen(0);
    await once(serverA, "listening");
    const portA = (serverA.address() as import("node:net").AddressInfo).port;

    // Create server B (where we'll subscribe)
    const appB = createServer(storeB as any, { 
      githubToken: "test-token",
      badgePubSub: sharedPubSub,
    });
    const serverB = appB.listen(0);
    await once(serverB, "listening");
    const portB = (serverB.address() as import("node:net").AddressInfo).port;

    // Connect a client to instance B and subscribe
    const clientB = new WebSocket(`ws://127.0.0.1:${portB}/api/ws`);
    const messagesB: any[] = [];
    clientB.on("message", (payload) => {
      messagesB.push(JSON.parse(payload.toString()));
    });
    await once(clientB, "open");
    clientB.send(JSON.stringify({ type: "subscribe", taskId: taskA.id }));

    // Give time for subscription to be established
    await new Promise((resolve) => setTimeout(resolve, 100));

    // Emit badge-changing task:updated on instance A (no local subscribers)
    const updatedTaskA = createTask({
      id: "FN-MULTI-001",
      prInfo: {
        url: "https://github.com/owner/repo/pull/1",
        number: 1,
        status: "merged",  // Status changed!
        title: "Merged PR",
        headBranch: "feature",
        baseBranch: "main",
        commentCount: 0,
        lastCheckedAt: "2026-03-30T12:00:00.000Z",
      },
      updatedAt: "2026-03-30T12:00:00.000Z",
    });
    (storeA as unknown as MockStore).task = updatedTaskA;
    (storeA as unknown as MockStore).emit("task:updated", updatedTaskA);

    // Wait for pub/sub propagation (InMemory uses setImmediate)
    await new Promise((resolve) => setTimeout(resolve, 200));

    // Cleanup first to avoid hanging
    clientB.close();
    await Promise.race([once(clientB, "close"), new Promise(r => setTimeout(r, 500))]);
    serverA.close();
    serverB.close();
    await Promise.race([once(serverA, "close"), new Promise(r => setTimeout(r, 500))]);
    await Promise.race([once(serverB, "close"), new Promise(r => setTimeout(r, 500))]);
    await sharedPubSub.dispose();

    // Verify instance B received the badge:updated message with merged status
    // Filter for the merged status message (might have received initial snapshot first)
    const mergedMessages = messagesB.filter(
      (m) => m.type === "badge:updated" && m.prInfo?.status === "merged"
    );
    expect(mergedMessages.length).toBeGreaterThanOrEqual(1);
    expect(mergedMessages[0]).toMatchObject({
      type: "badge:updated",
      taskId: taskA.id,
      prInfo: expect.objectContaining({ status: "merged" }),
    });
  }, 5000);

  websocketIntegrationTest("does not double-send badge updates to origin subscribers", async () => {
    // Create a shared pub/sub adapter
    const sharedPubSub: BadgePubSub = new InMemoryBadgePubSub();
    await sharedPubSub.start();

    const task = createTask({ id: "FN-ECHO-001" });
    const store = new MockStore(task);

    const app = createServer(store as any, { 
      githubToken: "test-token",
      badgePubSub: sharedPubSub,
    });
    const server = app.listen(0);
    await once(server, "listening");
    const port = (server.address() as import("node:net").AddressInfo).port;

    // Connect a client and subscribe
    const client = new WebSocket(`ws://127.0.0.1:${port}/api/ws`);
    const messages: any[] = [];
    client.on("message", (payload) => {
      messages.push(JSON.parse(payload.toString()));
    });
    await once(client, "open");
    client.send(JSON.stringify({ type: "subscribe", taskId: task.id }));

    // Wait for subscription
    await new Promise((resolve) => setTimeout(resolve, 100));

    // Emit task:updated on the same instance
    const updatedTask = createTask({
      id: "FN-ECHO-001",
      prInfo: {
        url: "https://github.com/owner/repo/pull/99",
        number: 99,
        status: "open",
        title: "New PR",
        headBranch: "feature",
        baseBranch: "main",
        commentCount: 0,
        lastCheckedAt: "2026-03-30T12:00:00.000Z",
      },
      updatedAt: "2026-03-30T12:00:00.000Z",
    });
    (store as unknown as MockStore).task = updatedTask;
    (store as unknown as MockStore).emit("task:updated", updatedTask);

    // Wait for any message delivery
    await new Promise((resolve) => setTimeout(resolve, 200));

    // Cleanup first
    client.close();
    await Promise.race([once(client, "close"), new Promise(r => setTimeout(r, 500))]);
    server.close();
    await Promise.race([once(server, "close"), new Promise(r => setTimeout(r, 500))]);
    await sharedPubSub.dispose();

    // Count badge:updated messages for this task with the new PR number
    const badgeMessages = messages.filter(
      (m) => m.type === "badge:updated" && m.taskId === task.id && m.prInfo?.number === 99
    );

    // With InMemoryBadgePubSub (which echoes messages for testing), we may receive:
    // 1. The local broadcast from onTaskUpdated
    // 2. The echoed message from pub/sub (InMemory doesn't filter by sourceId)
    // 
    // In production with RedisBadgePubSub, only 1 message would be received because
    // the Redis adapter filters out messages from the same sourceId.
    //
    // We verify that at least one message is received (the local broadcast)
    expect(badgeMessages.length).toBeGreaterThanOrEqual(1);
    expect(badgeMessages[0].prInfo.number).toBe(99);
  }, 5000);

  websocketIntegrationTest("sends cached badge snapshot to late subscribers after remote update", async () => {
    // Create a shared pub/sub adapter
    const sharedPubSub: BadgePubSub = new InMemoryBadgePubSub();
    await sharedPubSub.start();

    const task = createTask({ 
      id: "FN-LATE-001",
      prInfo: {
        url: "https://github.com/owner/repo/pull/1",
        number: 1,
        status: "open",
        title: "Original",
        headBranch: "feature",
        baseBranch: "main",
        commentCount: 0,
        lastCheckedAt: "2026-03-30T00:00:00.000Z",
      },
    });
    
    const storeA = new MockStore(task);
    const storeB = new MockStore(task);

    // Create both instances
    const appA = createServer(storeA as any, { 
      githubToken: "test-token",
      badgePubSub: sharedPubSub,
    });
    const serverA = appA.listen(0);
    await once(serverA, "listening");
    const portA = (serverA.address() as import("node:net").AddressInfo).port;

    const appB = createServer(storeB as any, { 
      githubToken: "test-token",
      badgePubSub: sharedPubSub,
    });
    const serverB = appB.listen(0);
    await once(serverB, "listening");
    const portB = (serverB.address() as import("node:net").AddressInfo).port;

    // First, emit an update on instance A (no subscribers)
    const updatedTask = createTask({
      id: "FN-LATE-001",
      prInfo: {
        url: "https://github.com/owner/repo/pull/1",
        number: 1,
        status: "merged",  // Changed!
        title: "Merged",
        headBranch: "feature",
        baseBranch: "main",
        commentCount: 0,
        lastCheckedAt: "2026-03-30T12:00:00.000Z",
      },
      updatedAt: "2026-03-30T12:00:00.000Z",
    });
    (storeA as unknown as MockStore).task = updatedTask;
    (storeA as unknown as MockStore).emit("task:updated", updatedTask);

    // Wait for pub/sub to propagate
    await new Promise((resolve) => setTimeout(resolve, 100));

    // Now connect a NEW client to instance B and subscribe
    const clientB = new WebSocket(`ws://127.0.0.1:${portB}/api/ws`);
    const messages: any[] = [];
    clientB.on("message", (payload) => {
      messages.push(JSON.parse(payload.toString()));
    });
    await once(clientB, "open");
    clientB.send(JSON.stringify({ type: "subscribe", taskId: task.id }));

    // Wait for late subscription replay
    await new Promise((resolve) => setTimeout(resolve, 200));

    // Cleanup
    clientB.close();
    await Promise.race([once(clientB, "close"), new Promise(r => setTimeout(r, 500))]);
    serverA.close();
    serverB.close();
    await Promise.race([once(serverA, "close"), new Promise(r => setTimeout(r, 500))]);
    await Promise.race([once(serverB, "close"), new Promise(r => setTimeout(r, 500))]);
    await sharedPubSub.dispose();

    // Verify the client received the cached "merged" snapshot (not stale "open")
    const badgeMessage = messages.find(m => m.type === "badge:updated");
    expect(badgeMessage).toBeDefined();
    expect(badgeMessage.prInfo.status).toBe("merged");
  }, 5000);
});