claude-code/src/utils/__tests__/collapseTeammateShutdowns.test.ts

import { describe, expect, test } from "bun:test";
import { collapseTeammateShutdowns } from "../collapseTeammateShutdowns";

function makeShutdownMsg(uuid = "1"): any {
  return {
    type: "attachment",
    uuid,
    timestamp: Date.now(),
    attachment: {
      type: "task_status",
      taskType: "in_process_teammate",
      status: "completed",
    },
  };
}

function makeNonShutdownMsg(): any {
  return { type: "user", message: { content: "hello" } };
}

describe("collapseTeammateShutdowns", () => {
  test("returns same messages when no teammate shutdowns", () => {
    const msgs = [makeNonShutdownMsg(), makeNonShutdownMsg()];
    expect(collapseTeammateShutdowns(msgs)).toEqual(msgs);
  });

  test("leaves single shutdown message unchanged", () => {
    const msgs = [makeShutdownMsg()];
    const result = collapseTeammateShutdowns(msgs);
    expect(result).toHaveLength(1);
    expect(result[0]).toEqual(msgs[0]);
  });

  test("collapses consecutive shutdown messages into batch", () => {
    const msgs = [makeShutdownMsg("1"), makeShutdownMsg("2")];
    const result = collapseTeammateShutdowns(msgs);
    expect(result).toHaveLength(1);
    expect(result[0].attachment.type).toBe("teammate_shutdown_batch");
  });

  test("batch attachment has correct count", () => {
    const msgs = [makeShutdownMsg("1"), makeShutdownMsg("2"), makeShutdownMsg("3")];
    const result = collapseTeammateShutdowns(msgs);
    expect(result[0].attachment.count).toBe(3);
  });

  test("does not collapse non-consecutive shutdowns", () => {
    const msgs = [makeShutdownMsg("1"), makeNonShutdownMsg(), makeShutdownMsg("2")];
    const result = collapseTeammateShutdowns(msgs);
    expect(result).toHaveLength(3);
    expect(result[0].attachment.type).toBe("task_status");
    expect(result[2].attachment.type).toBe("task_status");
  });

  test("preserves non-shutdown messages between shutdowns", () => {
    const msgs = [makeShutdownMsg("1"), makeNonShutdownMsg(), makeShutdownMsg("2")];
    const result = collapseTeammateShutdowns(msgs);
    expect(result[1]).toEqual(makeNonShutdownMsg());
  });

  test("handles empty array", () => {
    expect(collapseTeammateShutdowns([])).toEqual([]);
  });

  test("handles mixed message types", () => {
    const msgs = [makeNonShutdownMsg(), makeShutdownMsg("1"), makeShutdownMsg("2"), makeNonShutdownMsg()];
    const result = collapseTeammateShutdowns(msgs);
    expect(result).toHaveLength(3);
    expect(result[1].attachment.type).toBe("teammate_shutdown_batch");
  });

  test("collapses more than 2 consecutive shutdowns", () => {
    const msgs = Array.from({ length: 5 }, (_, i) => makeShutdownMsg(String(i)));
    const result = collapseTeammateShutdowns(msgs);
    expect(result).toHaveLength(1);
    expect(result[0].attachment.count).toBe(5);
  });

  test("non-teammate task_status messages are not collapsed", () => {
    const nonTeammate: any = {
      type: "attachment",
      uuid: "x",
      timestamp: Date.now(),
      attachment: {
        type: "task_status",
        taskType: "subagent",
        status: "completed",
      },
    };
    const msgs = [nonTeammate, { ...nonTeammate, uuid: "y" }];
    const result = collapseTeammateShutdowns(msgs);
    expect(result).toHaveLength(2);
  });
});