AgentHub/client/src/components/Sagi/__tests__/brain.test.ts

/**
 * @file brain.test.ts
 * @description Unit tests for the Sagi state brain — initial state, mood transitions driven by dashboard events, and status derivation.
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

import { describe, it, expect } from "vitest";
import {
  initialSagiState,
  reduceSagi,
  deriveMood,
  statusOf,
  clearErrors,
  seedSessions,
  HAPPY_MS,
  WORRIED_MS,
  STUCK_MS,
  SLEEP_MS,
  type SagiState,
} from "../brain";
import type {
  WSMessage,
  Session,
  Agent,
  DashboardEvent,
  RunStatusPayload,
} from "../../../lib/types";

const T0 = 1_000_000;

function sessionMsg(id: string, status: Session["status"], ts = T0): WSMessage {
  return { type: "session_updated", data: { id, status } as Session, timestamp: String(ts) };
}
function agentMsg(status: Agent["status"]): WSMessage {
  return { type: "agent_updated", data: { status } as Agent, timestamp: String(T0) };
}
function agentCreatedMsg(type: Agent["type"], status: Agent["status"] = "working"): WSMessage {
  return { type: "agent_created", data: { type, status } as Agent, timestamp: String(T0) };
}
function waitingMsg(id: string): WSMessage {
  return {
    type: "session_updated",
    data: { id, status: "active", awaiting_input_since: "2026-05-29T00:00:00Z" } as Session,
    timestamp: String(T0),
  };
}
function eventMsg(event_type: string): WSMessage {
  return { type: "new_event", data: { event_type } as DashboardEvent, timestamp: String(T0) };
}
function runStatusMsg(d: Partial<RunStatusPayload>): WSMessage {
  return { type: "run_status", data: d as RunStatusPayload, timestamp: String(T0) };
}

describe("deriveMood priority", () => {
  it("disconnected outranks everything", () => {
    const s: SagiState = { ...initialSagiState(T0), connected: false, worriedUntil: T0 + 9999 };
    expect(deriveMood(s, T0)).toBe("disconnected");
  });

  it("worried outranks stuck", () => {
    let s = initialSagiState(T0);
    ({ state: s } = reduceSagi(s, sessionMsg("a", "active"), T0)); // live
    s = { ...s, lastActivityAt: T0 - STUCK_MS - 1, worriedUntil: T0 + 100 };
    expect(deriveMood(s, T0)).toBe("worried");
  });

  it("stuck when a live session goes silent past STUCK_MS", () => {
    let s = initialSagiState(T0);
    ({ state: s } = reduceSagi(s, sessionMsg("a", "active"), T0));
    expect(deriveMood(s, T0 + STUCK_MS + 1)).toBe("stuck");
  });

  it("happy is transient then falls back to idle when nothing live", () => {
    let s = initialSagiState(T0);
    ({ state: s } = reduceSagi(s, sessionMsg("a", "active"), T0));
    ({ state: s } = reduceSagi(s, sessionMsg("a", "completed"), T0)); // no longer live
    expect(deriveMood(s, T0 + 10)).toBe("happy");
    expect(deriveMood(s, T0 + HAPPY_MS + 1)).toBe("idle");
  });

  it("thinking shows when set and nothing higher applies", () => {
    const s = { ...initialSagiState(T0), thinking: true };
    expect(deriveMood(s, T0)).toBe("thinking");
  });

  it("watching when a session is live and recent", () => {
    let s = initialSagiState(T0);
    ({ state: s } = reduceSagi(s, sessionMsg("a", "active"), T0));
    expect(deriveMood(s, T0 + 1000)).toBe("watching");
  });

  it("sleeping after SLEEP_MS of no activity and nothing live", () => {
    const s = initialSagiState(T0);
    expect(deriveMood(s, T0 + SLEEP_MS + 1)).toBe("sleeping");
  });

  it("idle by default", () => {
    expect(deriveMood(initialSagiState(T0), T0)).toBe("idle");
  });
});

describe("reduceSagi counts and pulses", () => {
  it("tracks live count accurately across transitions", () => {
    let s = initialSagiState(T0);
    ({ state: s } = reduceSagi(s, sessionMsg("a", "active"), T0));
    ({ state: s } = reduceSagi(s, sessionMsg("b", "active"), T0));
    expect(statusOf(s).liveCount).toBe(2);
    ({ state: s } = reduceSagi(s, sessionMsg("a", "completed"), T0));
    expect(statusOf(s).liveCount).toBe(1);
  });

  it("counts errored sessions and emits error pulse", () => {
    let s = initialSagiState(T0);
    const r = reduceSagi(s, sessionMsg("a", "error"), T0);
    s = r.state;
    expect(r.pulse).toBe("error");
    expect(statusOf(s).errorCount).toBe(1);
    expect(deriveMood(s, T0)).toBe("worried");
  });

  it("session_start pulse only on first active transition", () => {
    let s = initialSagiState(T0);
    const r1 = reduceSagi(s, sessionMsg("a", "active"), T0);
    expect(r1.pulse).toBe("session_start");
    const r2 = reduceSagi(r1.state, sessionMsg("a", "active"), T0);
    expect(r2.pulse).toBe(null);
  });

  it("session_done pulse only when the session was tracked", () => {
    let s = initialSagiState(T0);
    const untracked = reduceSagi(s, sessionMsg("ghost", "completed"), T0);
    expect(untracked.pulse).toBe(null);
    ({ state: s } = reduceSagi(s, sessionMsg("a", "active"), T0));
    const done = reduceSagi(s, sessionMsg("a", "completed"), T0);
    expect(done.pulse).toBe("session_done");
  });

  it("agent error triggers worried via pulse", () => {
    const r = reduceSagi(initialSagiState(T0), agentMsg("error"), T0);
    expect(r.pulse).toBe("error");
    expect(deriveMood(r.state, T0)).toBe("worried");
  });

  it("a newly created subagent emits subagent_spawn, a main agent does not", () => {
    expect(reduceSagi(initialSagiState(T0), agentCreatedMsg("subagent"), T0).pulse).toBe(
      "subagent_spawn"
    );
    expect(reduceSagi(initialSagiState(T0), agentCreatedMsg("main"), T0).pulse).toBe(null);
  });

  it("waiting transition emits a waiting pulse once and still counts as live", () => {
    let s = initialSagiState(T0);
    ({ state: s } = reduceSagi(s, sessionMsg("a", "active"), T0)); // active
    const first = reduceSagi(s, waitingMsg("a"), T0);
    expect(first.pulse).toBe("waiting");
    expect(statusOf(first.state).liveCount).toBe(1);
    // A repeat waiting update does not re-announce.
    const second = reduceSagi(first.state, waitingMsg("a"), T0);
    expect(second.pulse).toBe(null);
  });

  it("failure event types set worried, normal events do not", () => {
    const fail = reduceSagi(initialSagiState(T0), eventMsg("toolError"), T0);
    expect(fail.pulse).toBe("error");
    const ok = reduceSagi(initialSagiState(T0), eventMsg("postToolUse"), T0);
    expect(ok.pulse).toBe(null);
    expect(ok.state.worriedUntil).toBe(0);
  });

  it("run_status updates activity timestamp only (no exit code available)", () => {
    const running = reduceSagi(initialSagiState(T0), runStatusMsg({ status: "running" }), T0);
    expect(running.pulse).toBe(null);
    const gone = reduceSagi(initialSagiState(T0), runStatusMsg({ status: "gone" }), T0);
    expect(gone.pulse).toBe(null);
  });

  it("any handled message refreshes lastActivityAt", () => {
    const s = { ...initialSagiState(T0), lastActivityAt: T0 - 99999 };
    const { state } = reduceSagi(s, eventMsg("postToolUse"), T0 + 5);
    expect(state.lastActivityAt).toBe(T0 + 5);
  });

  it("ignores unrelated message types without mutating", () => {
    const s = initialSagiState(T0);
    const r = reduceSagi(s, { type: "import.progress", data: {} as never, timestamp: "x" }, T0);
    expect(r.state).toBe(s);
    expect(r.pulse).toBe(null);
  });
});

describe("seedSessions", () => {
  it("hydrates live/waiting/errored counts from a REST snapshot", () => {
    const s = seedSessions(
      initialSagiState(T0),
      [
        { id: "a", status: "active" },
        { id: "b", status: "active", awaiting_input_since: "2026-05-29T00:00:00Z" },
        { id: "c", status: "error" },
        { id: "d", status: "completed" }, // ignored
      ],
      T0
    );
    const st = statusOf(s);
    expect(st.liveCount).toBe(2); // a + b
    expect(st.waitingCount).toBe(1); // b
    expect(st.errorCount).toBe(1); // c
  });
});

describe("statusOf waiting", () => {
  it("counts a waiting session as both live and waiting", () => {
    let s = initialSagiState(T0);
    ({ state: s } = reduceSagi(s, waitingMsg("a"), T0));
    expect(statusOf(s)).toMatchObject({ liveCount: 1, waitingCount: 1, errorCount: 0 });
  });
});

describe("clearErrors", () => {
  it("drops errored sessions but keeps active ones", () => {
    let s = initialSagiState(T0);
    ({ state: s } = reduceSagi(s, sessionMsg("a", "active"), T0));
    ({ state: s } = reduceSagi(s, sessionMsg("b", "error"), T0));
    s = clearErrors(s);
    expect(statusOf(s).errorCount).toBe(0);
    expect(statusOf(s).liveCount).toBe(1);
    expect(s.worriedUntil).toBe(0);
  });

  it("keeps a cleared session quiet while it stays errored", () => {
    let s = initialSagiState(T0);
    ({ state: s } = reduceSagi(s, sessionMsg("b", "error"), T0));
    s = clearErrors(s);
    const again = reduceSagi(s, sessionMsg("b", "error", T0 + 1), T0 + 1);
    expect(again.pulse).toBeNull();
    expect(statusOf(again.state).errorCount).toBe(0);
    expect(again.state.worriedUntil).toBe(0);
    const ev = reduceSagi(
      again.state,
      { type: "new_event", data: { session_id: "b", event_type: "toolError" }, timestamp: "" } as WSMessage,
      T0 + 2
    );
    expect(ev.pulse).toBeNull();
    expect(seedSessions(s, [{ id: "b", status: "error" }], T0).sessions.b).toBeUndefined();
  });

  it("alerts again once a cleared session recovers and fails anew", () => {
    let s = initialSagiState(T0);
    ({ state: s } = reduceSagi(s, sessionMsg("b", "error"), T0));
    s = clearErrors(s);
    ({ state: s } = reduceSagi(s, sessionMsg("b", "active"), T0));
    expect(s.dismissed).not.toHaveProperty("b");
    const r = reduceSagi(s, sessionMsg("b", "error"), T0);
    expect(r.pulse).toBe("error");
  });
});

describe("error alerts fire on the transition only", () => {
  it("does not re-alert on each update of a still-errored session", () => {
    let s = initialSagiState(T0);
    ({ state: s } = reduceSagi(s, sessionMsg("a", "error"), T0));
    const r = reduceSagi(s, sessionMsg("a", "error"), T0 + 10_000);
    expect(r.pulse).toBeNull();
    expect(statusOf(r.state).errorCount).toBe(1);
  });
});

// Reference the imported constant so it is exercised and tsc-clean.
it("WORRIED_MS is a positive window", () => {
  expect(WORRIED_MS).toBeGreaterThan(0);
});