kandev/apps/web/hooks/domains/session/use-shell-command-output.test.ts

import { act, renderHook } from "@testing-library/react";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import {
  fetchShellCommandOutput,
  type ShellCommandOutputSnapshot,
} from "@/lib/api/domains/session-api";
import { useShellCommandOutput } from "./use-shell-command-output";

vi.mock("@/lib/api/domains/session-api", () => ({
  fetchShellCommandOutput: vi.fn(),
}));

const fetchOutputMock = vi.mocked(fetchShellCommandOutput);
let messageSequence = 0;
let messageId = "";

function snapshot(
  status: string,
  stdout = "",
  output: ShellCommandOutputSnapshot["output"] = {},
): ShellCommandOutputSnapshot {
  return {
    message_id: messageId,
    status,
    updated_at: "2026-07-16T12:00:00Z",
    output: { ...output, ...(stdout ? { stdout } : {}) },
  };
}

function deferred<T>() {
  let resolve!: (value: T) => void;
  let reject!: (reason?: unknown) => void;
  const promise = new Promise<T>((resolvePromise, rejectPromise) => {
    resolve = resolvePromise;
    reject = rejectPromise;
  });
  return { promise, resolve, reject };
}

async function flushPromises() {
  await act(async () => {
    await Promise.resolve();
    await Promise.resolve();
  });
}

beforeEach(() => {
  vi.useFakeTimers();
  fetchOutputMock.mockReset();
  messageSequence += 1;
  messageId = `message-${messageSequence}`;
});

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

describe("useShellCommandOutput fetching", () => {
  it("does not fetch while collapsed and fetches one terminal snapshot immediately on open", async () => {
    fetchOutputMock.mockResolvedValue(snapshot("complete", "done", { exit_code: 0 }));
    const { result, rerender } = renderHook(
      ({ isOpen }) =>
        useShellCommandOutput({
          sessionId: "session-1",
          messageId,
          isOpen,
          messageStatus: "complete",
        }),
      { initialProps: { isOpen: false } },
    );

    expect(fetchOutputMock).not.toHaveBeenCalled();
    rerender({ isOpen: true });
    expect(result.current.isLoading).toBe(true);
    await flushPromises();

    expect(result.current.snapshot?.output.stdout).toBe("done");
    expect(fetchOutputMock).toHaveBeenCalledTimes(1);
    await act(async () => vi.advanceTimersByTime(10_000));
    expect(fetchOutputMock).toHaveBeenCalledTimes(1);
  });

  it("polls running output after settlement without overlapping requests", async () => {
    const second = deferred<ShellCommandOutputSnapshot>();
    fetchOutputMock
      .mockResolvedValueOnce(snapshot("running", "first"))
      .mockReturnValueOnce(second.promise)
      .mockResolvedValueOnce(snapshot("complete", "unexpected"));
    const { result } = renderHook(() =>
      useShellCommandOutput({
        sessionId: "session-1",
        messageId,
        isOpen: true,
        messageStatus: "running",
      }),
    );
    await flushPromises();

    expect(result.current.snapshot?.output.stdout).toBe("first");
    await act(async () => vi.advanceTimersByTime(1_000));
    expect(fetchOutputMock).toHaveBeenCalledTimes(2);
    await act(async () => vi.advanceTimersByTime(10_000));
    expect(fetchOutputMock).toHaveBeenCalledTimes(2);

    second.resolve(snapshot("complete", "finished", { exit_code: 0 }));
    await flushPromises();
    expect(result.current.snapshot?.output.stdout).toBe("finished");
    await act(async () => vi.advanceTimersByTime(10_000));
    expect(fetchOutputMock).toHaveBeenCalledTimes(2);
  });

  it("polls pending output while the projected message remains active", async () => {
    fetchOutputMock
      .mockResolvedValueOnce(snapshot("pending", "queued"))
      .mockResolvedValueOnce(snapshot("complete", "done", { exit_code: 0 }));
    const { result } = renderHook(() =>
      useShellCommandOutput({
        sessionId: "session-1",
        messageId,
        isOpen: true,
        messageStatus: "pending",
      }),
    );
    await flushPromises();

    expect(result.current.snapshot?.output.stdout).toBe("queued");
    await act(async () => vi.advanceTimersByTime(1_000));
    await flushPromises();
    expect(fetchOutputMock).toHaveBeenCalledTimes(2);
  });

  it("retains the latest snapshot and caps consecutive failure backoff at five seconds", async () => {
    fetchOutputMock
      .mockResolvedValueOnce(snapshot("running", "retained"))
      .mockRejectedValue(new Error("unavailable"));
    const { result } = renderHook(() =>
      useShellCommandOutput({
        sessionId: "session-1",
        messageId,
        isOpen: true,
        messageStatus: "running",
      }),
    );
    await flushPromises();

    const advanceToNextFailure = async (delay: number, expectedCalls: number) => {
      await act(async () => vi.advanceTimersByTime(delay));
      await flushPromises();
      expect(fetchOutputMock).toHaveBeenCalledTimes(expectedCalls);
      expect(result.current.snapshot?.output.stdout).toBe("retained");
      expect(result.current.error?.message).toBe("unavailable");
    };

    await advanceToNextFailure(1_000, 2);
    await advanceToNextFailure(1_000, 3);
    await advanceToNextFailure(2_000, 4);
    await advanceToNextFailure(4_000, 5);
    await advanceToNextFailure(5_000, 6);
    await advanceToNextFailure(5_000, 7);
  });
});

describe("useShellCommandOutput shared cache", () => {
  it("deduplicates a concurrent request across hook instances", async () => {
    const pending = deferred<ShellCommandOutputSnapshot>();
    fetchOutputMock.mockReturnValue(pending.promise);

    const { result } = renderHook(() => {
      const first = useShellCommandOutput({
        sessionId: "session-1",
        messageId,
        isOpen: true,
        messageStatus: "complete",
      });
      const second = useShellCommandOutput({
        sessionId: "session-1",
        messageId,
        isOpen: true,
        messageStatus: "complete",
      });
      return { first, second };
    });

    expect(fetchOutputMock).toHaveBeenCalledTimes(1);
    pending.resolve(snapshot("complete", "shared", { exit_code: 0 }));
    await flushPromises();
    expect(result.current.first.snapshot?.output.stdout).toBe("shared");
    expect(result.current.second.snapshot?.output.stdout).toBe("shared");
  });

  it("reuses a terminal payload after the renderer is closed and reopened", async () => {
    fetchOutputMock.mockResolvedValue(snapshot("complete", "cached", { exit_code: 0 }));
    const first = renderHook(() =>
      useShellCommandOutput({
        sessionId: "session-1",
        messageId,
        isOpen: true,
        messageStatus: "complete",
      }),
    );
    await flushPromises();
    expect(first.result.current.snapshot?.output.stdout).toBe("cached");
    first.unmount();

    const reopened = renderHook(() =>
      useShellCommandOutput({
        sessionId: "session-1",
        messageId,
        isOpen: true,
        messageStatus: "complete",
      }),
    );
    await flushPromises();

    expect(reopened.result.current.snapshot?.output.stdout).toBe("cached");
    expect(fetchOutputMock).toHaveBeenCalledTimes(1);
  });

  it("replaces a shared running request once when separate renderers become terminal", async () => {
    const finalOutput = "final shared output";
    const running = deferred<ShellCommandOutputSnapshot>();
    const terminal = deferred<ShellCommandOutputSnapshot>();
    const signals: AbortSignal[] = [];
    fetchOutputMock.mockImplementation((_sessionId, _messageId, options) => {
      const signal = options?.init?.signal;
      if (signal) signals.push(signal);
      return signals.length === 1 ? running.promise : terminal.promise;
    });

    const renderOutput = () =>
      renderHook(
        ({ messageStatus }) =>
          useShellCommandOutput({
            sessionId: "session-1",
            messageId,
            isOpen: true,
            messageStatus,
          }),
        { initialProps: { messageStatus: "running" } },
      );
    const first = renderOutput();
    const second = renderOutput();
    expect(fetchOutputMock).toHaveBeenCalledTimes(1);

    first.rerender({ messageStatus: "complete" });
    second.rerender({ messageStatus: "complete" });

    expect(signals[0]?.aborted).toBe(true);
    expect(fetchOutputMock).toHaveBeenCalledTimes(2);
    terminal.resolve(snapshot("complete", finalOutput, { exit_code: 0 }));
    await flushPromises();
    expect(first.result.current.snapshot?.output.stdout).toBe(finalOutput);
    expect(second.result.current.snapshot?.output.stdout).toBe(finalOutput);

    running.resolve(snapshot("running", "stale shared output"));
    await flushPromises();
    first.unmount();
    second.unmount();

    const reopened = renderOutput();
    await flushPromises();
    expect(reopened.result.current.snapshot?.output.stdout).toBe(finalOutput);
    expect(fetchOutputMock).toHaveBeenCalledTimes(2);
  });
});

describe("useShellCommandOutput cleanup", () => {
  it("aborts and ignores stale output after collapse", async () => {
    const pending = deferred<ShellCommandOutputSnapshot>();
    let requestSignal: AbortSignal | undefined;
    fetchOutputMock.mockImplementation((_sessionId, _messageId, options) => {
      requestSignal = options?.init?.signal ?? undefined;
      return pending.promise;
    });
    const { result, rerender } = renderHook(
      ({ isOpen }) =>
        useShellCommandOutput({
          sessionId: "session-1",
          messageId,
          isOpen,
          messageStatus: "running",
        }),
      { initialProps: { isOpen: true } },
    );

    expect(requestSignal?.aborted).toBe(false);
    rerender({ isOpen: false });
    expect(requestSignal?.aborted).toBe(true);
    pending.resolve(snapshot("running", "stale"));
    await flushPromises();
    expect(result.current.snapshot).toBeNull();
  });

  it("replaces in-flight work with one final fetch when the projected message becomes terminal", async () => {
    const signals: AbortSignal[] = [];
    const finalSnapshot = deferred<ShellCommandOutputSnapshot>();
    fetchOutputMock.mockImplementation((_sessionId, _messageId, options) => {
      const signal = options?.init?.signal;
      if (signal) signals.push(signal);
      if (signals.length === 2) return finalSnapshot.promise;
      return new Promise<ShellCommandOutputSnapshot>(() => undefined);
    });
    const projectedTerminal = renderHook(
      ({ messageStatus }) =>
        useShellCommandOutput({
          sessionId: "session-1",
          messageId,
          isOpen: true,
          messageStatus,
        }),
      { initialProps: { messageStatus: "running" } },
    );

    projectedTerminal.rerender({ messageStatus: "complete" });
    expect(signals[0]?.aborted).toBe(true);
    expect(signals).toHaveLength(2);
    expect(signals[1]?.aborted).toBe(false);
    finalSnapshot.resolve(snapshot("complete", "final output", { exit_code: 0 }));
    await flushPromises();
    expect(projectedTerminal.result.current.snapshot?.output.stdout).toBe("final output");
    await act(async () => vi.advanceTimersByTime(10_000));
    expect(fetchOutputMock).toHaveBeenCalledTimes(2);
    projectedTerminal.unmount();
  });
});

describe("useShellCommandOutput terminal lifecycle", () => {
  it("aborts in-flight work on unmount", () => {
    const signals: AbortSignal[] = [];
    fetchOutputMock.mockImplementation((_sessionId, _messageId, options) => {
      const signal = options?.init?.signal;
      if (signal) signals.push(signal);
      return new Promise<ShellCommandOutputSnapshot>(() => undefined);
    });
    const unmounted = renderHook(() =>
      useShellCommandOutput({
        sessionId: "session-1",
        messageId,
        isOpen: true,
        messageStatus: "running",
      }),
    );
    expect(signals).toHaveLength(1);
    unmounted.unmount();
    expect(signals[0]?.aborted).toBe(true);
  });

  it.each(["completed", "success", "failed"])(
    "treats the projected %s alias as terminal after the final snapshot",
    async (terminalStatus) => {
      fetchOutputMock
        .mockResolvedValueOnce(snapshot("running", "partial"))
        .mockResolvedValueOnce(snapshot("running", "finalizing"));
      const view = renderHook(
        ({ messageStatus }) =>
          useShellCommandOutput({
            sessionId: "session-1",
            messageId,
            isOpen: true,
            messageStatus,
          }),
        { initialProps: { messageStatus: "running" } },
      );
      await flushPromises();

      view.rerender({ messageStatus: terminalStatus });
      await flushPromises();
      expect(fetchOutputMock).toHaveBeenCalledTimes(2);
      await act(async () => vi.advanceTimersByTime(10_000));
      expect(fetchOutputMock).toHaveBeenCalledTimes(2);
    },
  );
});

it("drops cached output and stops retrying when the server reports removal", async () => {
  const { ApiError } = await import("@/lib/api/client");
  fetchOutputMock
    .mockResolvedValueOnce(snapshot("running", "deleted payload"))
    .mockRejectedValue(new ApiError("gone", 410, { code: "tool_payload_removed" }));
  const { result } = renderHook(() =>
    useShellCommandOutput({
      sessionId: "removed-session-1",
      messageId: "message-1",
      isOpen: true,
      messageStatus: "running",
    }),
  );
  await flushPromises();
  await act(async () => vi.advanceTimersByTime(1000));
  expect(result.current.snapshot).toBeNull();
  await act(async () => vi.advanceTimersByTime(10000));
  expect(fetchOutputMock).toHaveBeenCalledTimes(2);
});