stack/apps/web/src/hooks/__tests__/useAgentStream.test.ts

/**
 * @file useAgentStream.test.ts
 * @description Unit tests for the useAgentStream hook
 *
 * Tests cover:
 * - SSE event parsing (agent:spawned, agent:output, agent:completed, agent:error)
 * - Agent lifecycle state transitions
 * - Auto-reconnect behavior on connection loss
 * - Cleanup on unmount
 * - Dismiss agent
 */

import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { renderHook, act } from "@testing-library/react";
import { useAgentStream } from "../useAgentStream";

// ==========================================
// Mock EventSource
// ==========================================

interface MockEventSourceInstance {
  url: string;
  onopen: (() => void) | null;
  onerror: ((event: Event) => void) | null;
  onmessage: ((event: MessageEvent) => void) | null;
  close: ReturnType<typeof vi.fn>;
  addEventListener: ReturnType<typeof vi.fn>;
  dispatchEvent: (type: string, data: string) => void;
  _listeners: Map<string, ((event: MessageEvent<string>) => void)[]>;
  readyState: number;
}

let mockEventSourceInstances: MockEventSourceInstance[] = [];

const MockEventSource = vi.fn(function (this: MockEventSourceInstance, url: string): void {
  this.url = url;
  this.onopen = null;
  this.onerror = null;
  this.onmessage = null;
  this.close = vi.fn();
  this.readyState = 0;
  this._listeners = new Map();

  this.addEventListener = vi.fn(
    (type: string, handler: (event: MessageEvent<string>) => void): void => {
      if (!this._listeners.has(type)) {
        this._listeners.set(type, []);
      }
      const list = this._listeners.get(type);
      if (list) list.push(handler);
    }
  );

  this.dispatchEvent = (type: string, data: string): void => {
    const handlers = this._listeners.get(type) ?? [];
    const event = new MessageEvent(type, { data });
    for (const handler of handlers) {
      handler(event);
    }
  };

  mockEventSourceInstances.push(this);
});

// Add static constants
Object.assign(MockEventSource, {
  CONNECTING: 0,
  OPEN: 1,
  CLOSED: 2,
});

vi.stubGlobal("EventSource", MockEventSource);

// ==========================================
// Helpers
// ==========================================

function getLatestES(): MockEventSourceInstance {
  const es = mockEventSourceInstances[mockEventSourceInstances.length - 1];
  if (!es) throw new Error("No EventSource instance created");
  return es;
}

function triggerOpen(): void {
  const es = getLatestES();
  if (es.onopen) es.onopen();
}

function triggerError(): void {
  const es = getLatestES();
  if (es.onerror) es.onerror(new Event("error"));
}

function emitEvent(type: string, data: unknown): void {
  const es = getLatestES();
  es.dispatchEvent(type, JSON.stringify(data));
}

// ==========================================
// Tests
// ==========================================

describe("useAgentStream", () => {
  beforeEach(() => {
    vi.clearAllMocks();
    vi.useFakeTimers();
    mockEventSourceInstances = [];
  });

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

  // ==========================================
  // Initialization
  // ==========================================

  describe("initialization", () => {
    it("creates an EventSource connecting to /api/orchestrator/events", () => {
      renderHook(() => useAgentStream());
      expect(MockEventSource).toHaveBeenCalledWith("/api/orchestrator/events");
    });

    it("starts with isConnected=false before onopen fires", () => {
      const { result } = renderHook(() => useAgentStream());
      expect(result.current.isConnected).toBe(false);
    });

    it("starts with an empty agents map", () => {
      const { result } = renderHook(() => useAgentStream());
      expect(result.current.agents.size).toBe(0);
    });

    it("sets isConnected=true when EventSource opens", () => {
      const { result } = renderHook(() => useAgentStream());
      act(() => {
        triggerOpen();
      });
      expect(result.current.isConnected).toBe(true);
    });

    it("clears connectionError when EventSource opens", () => {
      const { result } = renderHook(() => useAgentStream());

      // Trigger an error first to set connectionError
      act(() => {
        triggerError();
      });

      // Start a fresh reconnect and open it
      act(() => {
        vi.advanceTimersByTime(2000);
      });

      act(() => {
        triggerOpen();
      });

      expect(result.current.connectionError).toBeNull();
    });
  });

  // ==========================================
  // SSE event: agent:spawned
  // ==========================================

  describe("agent:spawned event", () => {
    it("adds an agent with status=spawning", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker", jobId: "job-abc" });
      });

      expect(result.current.agents.has("agent-1")).toBe(true);
      expect(result.current.agents.get("agent-1")?.status).toBe("spawning");
    });

    it("sets agentType from the type field", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "planner" });
      });

      expect(result.current.agents.get("agent-1")?.agentType).toBe("planner");
    });

    it("defaults agentType to 'agent' when type is missing", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-2" });
      });

      expect(result.current.agents.get("agent-2")?.agentType).toBe("agent");
    });

    it("stores jobId when present", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-3", type: "worker", jobId: "job-xyz" });
      });

      expect(result.current.agents.get("agent-3")?.jobId).toBe("job-xyz");
    });

    it("starts with empty outputLines", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
      });

      expect(result.current.agents.get("agent-1")?.outputLines).toEqual([]);
    });
  });

  // ==========================================
  // SSE event: agent:output
  // ==========================================

  describe("agent:output event", () => {
    it("appends output to the agent's outputLines", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
        emitEvent("agent:output", { agentId: "agent-1", data: "Hello world\n" });
      });

      expect(result.current.agents.get("agent-1")?.outputLines).toContain("Hello world\n");
    });

    it("transitions status from spawning to running on first output", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
        emitEvent("agent:output", { agentId: "agent-1", data: "Starting...\n" });
      });

      expect(result.current.agents.get("agent-1")?.status).toBe("running");
    });

    it("accumulates multiple output lines", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
        emitEvent("agent:output", { agentId: "agent-1", data: "Line 1\n" });
        emitEvent("agent:output", { agentId: "agent-1", data: "Line 2\n" });
        emitEvent("agent:output", { agentId: "agent-1", data: "Line 3\n" });
      });

      const lines = result.current.agents.get("agent-1")?.outputLines ?? [];
      expect(lines).toHaveLength(3);
      expect(lines[0]).toBe("Line 1\n");
      expect(lines[1]).toBe("Line 2\n");
      expect(lines[2]).toBe("Line 3\n");
    });

    it("creates a new agent entry if output arrives before spawned event", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:output", { agentId: "unknown-agent", data: "Surprise output\n" });
      });

      expect(result.current.agents.has("unknown-agent")).toBe(true);
      expect(result.current.agents.get("unknown-agent")?.status).toBe("running");
    });
  });

  // ==========================================
  // SSE event: agent:completed
  // ==========================================

  describe("agent:completed event", () => {
    it("sets status to completed", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
        emitEvent("agent:output", { agentId: "agent-1", data: "Working...\n" });
        emitEvent("agent:completed", { agentId: "agent-1", exitCode: 0 });
      });

      expect(result.current.agents.get("agent-1")?.status).toBe("completed");
    });

    it("stores the exitCode", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
        emitEvent("agent:completed", { agentId: "agent-1", exitCode: 42 });
      });

      expect(result.current.agents.get("agent-1")?.exitCode).toBe(42);
    });

    it("sets endedAt timestamp", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
        emitEvent("agent:completed", { agentId: "agent-1", exitCode: 0 });
      });

      expect(result.current.agents.get("agent-1")?.endedAt).toBeDefined();
    });

    it("ignores completed event for unknown agent", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:completed", { agentId: "ghost-agent", exitCode: 0 });
      });

      expect(result.current.agents.has("ghost-agent")).toBe(false);
    });
  });

  // ==========================================
  // SSE event: agent:error
  // ==========================================

  describe("agent:error event", () => {
    it("sets status to error", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
        emitEvent("agent:error", { agentId: "agent-1", error: "Out of memory" });
      });

      expect(result.current.agents.get("agent-1")?.status).toBe("error");
    });

    it("stores the error message", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
        emitEvent("agent:error", { agentId: "agent-1", error: "Segfault" });
      });

      expect(result.current.agents.get("agent-1")?.errorMessage).toBe("Segfault");
    });

    it("sets endedAt on error", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
        emitEvent("agent:error", { agentId: "agent-1", error: "Crash" });
      });

      expect(result.current.agents.get("agent-1")?.endedAt).toBeDefined();
    });

    it("ignores error event for unknown agent", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:error", { agentId: "ghost-agent", error: "Crash" });
      });

      expect(result.current.agents.has("ghost-agent")).toBe(false);
    });
  });

  // ==========================================
  // Reconnect behavior
  // ==========================================

  describe("auto-reconnect", () => {
    it("sets isConnected=false on error", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
      });
      act(() => {
        triggerError();
      });

      expect(result.current.isConnected).toBe(false);
    });

    it("sets connectionError on error", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        triggerError();
      });

      expect(result.current.connectionError).not.toBeNull();
    });

    it("creates a new EventSource after reconnect delay", () => {
      renderHook(() => useAgentStream());
      const initialCount = mockEventSourceInstances.length;

      act(() => {
        triggerOpen();
        triggerError();
      });

      act(() => {
        vi.advanceTimersByTime(1500); // initial backoff = 1000ms
      });

      expect(mockEventSourceInstances.length).toBeGreaterThan(initialCount);
    });

    it("closes the old EventSource before reconnecting", () => {
      renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        triggerError();
      });

      const closedInstance = mockEventSourceInstances[0];
      expect(closedInstance?.close).toHaveBeenCalled();
    });
  });

  // ==========================================
  // Cleanup on unmount
  // ==========================================

  describe("cleanup on unmount", () => {
    it("closes EventSource when the hook is unmounted", () => {
      const { unmount } = renderHook(() => useAgentStream());

      const es = getLatestES();
      unmount();

      expect(es.close).toHaveBeenCalled();
    });

    it("does not attempt to reconnect after unmount", () => {
      const { unmount } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        triggerError();
      });

      const countBeforeUnmount = mockEventSourceInstances.length;

      unmount();

      act(() => {
        vi.advanceTimersByTime(5000);
      });

      // No new instances created after unmount
      expect(mockEventSourceInstances.length).toBe(countBeforeUnmount);
    });
  });

  // ==========================================
  // Dismiss agent
  // ==========================================

  describe("dismissAgent", () => {
    it("removes the agent from the map", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
        emitEvent("agent:completed", { agentId: "agent-1", exitCode: 0 });
      });

      act(() => {
        result.current.dismissAgent("agent-1");
      });

      expect(result.current.agents.has("agent-1")).toBe(false);
    });

    it("is a no-op for unknown agentId", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        emitEvent("agent:spawned", { agentId: "agent-1", type: "worker" });
      });

      act(() => {
        result.current.dismissAgent("nonexistent-agent");
      });

      expect(result.current.agents.has("agent-1")).toBe(true);
    });
  });

  // ==========================================
  // Malformed event handling
  // ==========================================

  describe("malformed events", () => {
    it("ignores malformed JSON without throwing", () => {
      const { result } = renderHook(() => useAgentStream());

      act(() => {
        triggerOpen();
        // Dispatch raw bad JSON
        const es = getLatestES();
        es.dispatchEvent("agent:spawned", "NOT JSON {{{");
      });

      // Should not crash, agents map stays empty
      expect(result.current.agents.size).toBe(0);
    });
  });
});