MEGAINFRA-WEBAPP/src/app/api/chat/route.test.ts

import { beforeEach, describe, expect, it, vi } from "vitest";
import type { ChatProvider, RunTurnArgs } from "@/lib/chat";
import type { ChartSpec } from "@/lib/reports/chart-spec";
// Static import (vi.mock is hoisted above it): loading the route pulls in
// the provider SDKs, which is slow under a parallel run and must count
// as collection, not test time.
import { POST } from "./route";

const spec: ChartSpec = {
  id: "site-weekly",
  kind: "bar",
  title: "Income vs expenses",
  xKey: "week",
  series: [{ key: "expenses", label: "Expenses", format: "money" }],
  rows: [{ week: "2026-W38", expenses: 987654 }],
};

vi.mock("@/lib/agent/tools", () => ({
  tools: [
    {
      name: "fake_report",
      description: "A fake report tool",
      schema: {},
      handler: async () => ({ total: "R 5", charts: [spec] }),
    },
  ],
}));

/** First turn calls the tool; second turn finishes. Records what it saw. */
const seen: RunTurnArgs[] = [];
/** Set by a test to make the next runTurn throw, mid-stream. */
let throwOnNextTurn: string | null = null;
/** Set by a test to make the next runTurn request a specific tool call. */
let toolCallOnNextTurn: { name: string; input: Record<string, unknown> } | null = null;
const provider: ChatProvider = {
  name: "fake",
  model: "fake-1",
  async runTurn(args) {
    seen.push({ ...args, history: [...args.history] });
    if (throwOnNextTurn) {
      const message = throwOnNextTurn;
      throwOnNextTurn = null;
      args.onTextDelta("Partial");
      throw new Error(message);
    }
    if (toolCallOnNextTurn) {
      const call = toolCallOnNextTurn;
      toolCallOnNextTurn = null;
      return { text: "", toolCalls: [{ id: "t1", name: call.name, input: call.input }], done: false };
    }
    if (seen.length === 1) {
      return { text: "", toolCalls: [{ id: "t1", name: "fake_report", input: {} }], done: false };
    }
    args.onTextDelta("Done.");
    return { text: "Done.", toolCalls: [], done: true };
  },
};

vi.mock("@/lib/chat", async (importOriginal) => {
  const actual = await importOriginal<typeof import("@/lib/chat")>();
  return { ...actual, chatEnabled: () => true, getProvider: () => provider };
});

const resolveActorMock = vi.fn();
vi.mock("@/lib/agent/context", () => ({
  resolveActor: (...args: unknown[]) => resolveActorMock(...args),
}));

const createThreadMock = vi.fn();
const getThreadOwnerMock = vi.fn();
const loadThreadEntriesMock = vi.fn();
const appendMessageMock = vi.fn();
const getRecentActionIdsMock = vi.fn(async () => [] as string[]);
vi.mock("@/lib/chat/threads", () => ({
  autoTitle: (text: string) => text.slice(0, 10),
  createThread: (...args: unknown[]) => createThreadMock(...args),
  getThreadOwner: (...args: unknown[]) => getThreadOwnerMock(...args),
  loadThreadEntries: (...args: unknown[]) => loadThreadEntriesMock(...args),
  appendMessage: (...args: unknown[]) => appendMessageMock(...args),
  getRecentActionIds: () => getRecentActionIdsMock(),
}));

const describeActionOutcomeMock = vi.fn<(...args: unknown[]) => Promise<string | null>>(
  async () => null,
);
const proposeChangesMock = vi.fn();
vi.mock("@/lib/agent/gateway", async (importOriginal) => {
  const actual = await importOriginal<typeof import("@/lib/agent/gateway")>();
  return {
    ...actual,
    describeActionOutcome: (id: string) => describeActionOutcomeMock(id),
    proposeChanges: (...args: unknown[]) => proposeChangesMock(...args),
  };
});

const listMemoriesMock = vi.fn<(...args: unknown[]) => Promise<{ text: string }[]>>(
  async () => [],
);
vi.mock("@/lib/agent/memory", () => ({
  listMemories: (...args: unknown[]) => listMemoriesMock(...args),
  formatMemoriesForPrompt: (memories: { text: string }[]) =>
    memories.map((m) => `- ${m.text}`).join("\n"),
}));

describe("POST /api/chat", () => {
  it("streams chart specs to the browser and only ids/titles to the model", async () => {
    const res = await POST(
      new Request("http://test/api/chat", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ history: [{ role: "user", text: "report" }] }),
      }),
    );
    const body = await res.text();

    expect(body).toContain(`event: chart\ndata: ${JSON.stringify(spec)}`);

    const second = seen[1];
    const toolEntry = second.history.find((e) => e.role === "tool");
    expect(toolEntry && toolEntry.role === "tool").toBe(true);
    const content = toolEntry && toolEntry.role === "tool" ? toolEntry.results[0].content : "";
    expect(content).toContain("chartsAvailable");
    expect(content).toContain("R 5");
    expect(content).not.toContain("987654");
    expect(content).not.toContain("2026-W38");
  });
});

function threadedRequest(body: Record<string, unknown>, headers: Record<string, string> = { "x-ask-user": "u1" }) {
  return new Request("http://test/api/chat", {
    method: "POST",
    headers: { "Content-Type": "application/json", ...headers },
    body: JSON.stringify(body),
  });
}

describe("POST /api/chat — threaded", () => {
  beforeEach(() => {
    resolveActorMock.mockReset();
    createThreadMock.mockReset();
    getThreadOwnerMock.mockReset();
    loadThreadEntriesMock.mockReset();
    appendMessageMock.mockReset();
    resolveActorMock.mockResolvedValue({ id: "u1", name: "Dashboard" });
    loadThreadEntriesMock.mockResolvedValue([]);
    appendMessageMock.mockResolvedValue("m1");
    listMemoriesMock.mockClear();
    listMemoriesMock.mockResolvedValue([]);
    getRecentActionIdsMock.mockClear();
    getRecentActionIdsMock.mockResolvedValue([]);
    describeActionOutcomeMock.mockClear();
    describeActionOutcomeMock.mockResolvedValue(null);
  });

  it("embeds the actor's memories in the system prompt and offers the remember tool", async () => {
    createThreadMock.mockResolvedValue({ id: "t9", title: "Report" });
    listMemoriesMock.mockResolvedValue([{ text: "The Hills works 45-hour weeks" }]);
    await POST(threadedRequest({ message: "How many hours does The Hills work?" }));

    const last = seen[seen.length - 1];
    expect(last.system).toContain("The Hills works 45-hour weeks");
    expect(last.tools.map((t) => t.name)).toContain("remember");
  });

  it("rejects a request with no x-ask-user (via resolveActor)", async () => {
    resolveActorMock.mockResolvedValue(
      Response.json({ error: "Pick who is using Ask" }, { status: 400 }),
    );
    const res = await POST(threadedRequest({ message: "hi" }, {}));
    expect(res.status).toBe(400);
    expect(createThreadMock).not.toHaveBeenCalled();
  });

  it("creates a thread, streams a thread event, and creates the user row before the model runs, the assistant row after", async () => {
    createThreadMock.mockResolvedValue({ id: "t9", title: "Report" });
    appendMessageMock.mockResolvedValueOnce("msg-user").mockResolvedValueOnce("msg-assistant");
    const res = await POST(threadedRequest({ message: "Report on The Hills" }));
    const body = await res.text();

    expect(body).toContain(`event: thread\ndata: ${JSON.stringify({ id: "t9", title: "Report" })}`);
    expect(appendMessageMock).toHaveBeenCalledTimes(2);
    expect(appendMessageMock.mock.calls[0][2]).toMatchObject({ kind: "user", text: "Report on The Hills" });
    expect(appendMessageMock.mock.calls[1][2]).toMatchObject({ kind: "assistant" });
    // The persisted row's id rides on "done" so the browser knows which
    // message to ask GET /api/ask/messages/[id]/pdf for.
    expect(body).toContain('"messageId":"msg-assistant"');
  });

  it("404s a thread owned by someone else", async () => {
    getThreadOwnerMock.mockResolvedValue("someone-else");
    const res = await POST(threadedRequest({ threadId: "t1", message: "hi" }));
    expect(res.status).toBe(404);
    expect(appendMessageMock).not.toHaveBeenCalled();
  });

  it("404s a thread that doesn't exist", async () => {
    getThreadOwnerMock.mockResolvedValue(null);
    const res = await POST(threadedRequest({ threadId: "t1", message: "hi" }));
    expect(res.status).toBe(404);
  });

  it("persists a partial assistant row with tools marked failed when the model errors mid-stream", async () => {
    createThreadMock.mockResolvedValue({ id: "t9", title: "Report" });
    throwOnNextTurn = "boom";
    const res = await POST(threadedRequest({ message: "Report on The Hills" }));
    const body = await res.text();

    expect(body).toContain("event: error");
    expect(appendMessageMock).toHaveBeenCalledTimes(2);
    const assistantCall = appendMessageMock.mock.calls[1][2];
    expect(assistantCall.kind).toBe("assistant");
    expect(assistantCall.tools).toContain("error");
  });

  it("streams the full action card and persists only the actionId on the assistant row", async () => {
    createThreadMock.mockResolvedValue({ id: "t9", title: "Add a cost" });
    proposeChangesMock.mockResolvedValue({
      actionId: "log1",
      status: "awaiting_confirmation",
      problems: [],
      card: { actionId: "log1", summary: "Add a fixed cost", preview: [], problems: [], status: "pending" },
    });
    toolCallOnNextTurn = { name: "propose_changes", input: { summary: "Add a fixed cost", ops: [] } };

    const res = await POST(threadedRequest({ message: "Add a fixed cost of R1 on Cedar Gate" }));
    const body = await res.text();

    expect(body).toContain(`event: action\ndata: ${JSON.stringify({
      actionId: "log1",
      summary: "Add a fixed cost",
      preview: [],
      problems: [],
      status: "pending",
    })}`);
    // The model never sees the preview — only the minimal fields.
    const toolEntry = seen
      .flatMap((s) => s.history)
      .find((e) => e.role === "tool") as { role: "tool"; results: { content: string }[] } | undefined;
    const lastToolResult = toolEntry?.results[toolEntry.results.length - 1]?.content ?? "";
    expect(lastToolResult).not.toContain("preview");

    const assistantCall = appendMessageMock.mock.calls.find((c) => c[2].kind === "assistant")?.[2];
    expect(assistantCall.actionIds).toEqual(["log1"]);
  });

  it("streams a pdf event for make_pdf_report and persists the structured document on the assistant row", async () => {
    createThreadMock.mockResolvedValue({ id: "t9", title: "Fuel report" });
    appendMessageMock.mockResolvedValue("msg-1");
    const doc = {
      title: "Fuel at The Hills",
      subtitle: "September 2026",
      sections: [{ heading: "Overview", text: "Fuel spend was steady." }],
    };
    toolCallOnNextTurn = { name: "make_pdf_report", input: doc };

    const res = await POST(threadedRequest({ message: "Give me a PDF on fuel at The Hills for September" }));
    const body = await res.text();

    expect(body).toContain(`event: pdf\ndata: ${JSON.stringify(doc)}`);
    expect(body).toContain('"messageId":"msg-1"');
    const assistantCall = appendMessageMock.mock.calls.find((c) => c[2].kind === "assistant")?.[2];
    expect(assistantCall.parsed).toEqual({ pdf: doc });
  });

  it("prepends the prior action's outcome as a system note for the next turn", async () => {
    getThreadOwnerMock.mockResolvedValue("u1");
    getRecentActionIdsMock.mockResolvedValue(["log1"]);
    describeActionOutcomeMock.mockResolvedValue('Action log1 "Add a fixed cost" was applied.');

    await POST(threadedRequest({ threadId: "t1", message: "thanks" }));

    const userCall = appendMessageMock.mock.calls.find((c) => c[2].kind === "user")?.[2];
    // The stored row keeps the user's own words — the note is only added
    // to what the model reads, not to what's persisted.
    expect(userCall.text).toBe("thanks");
    const modelHistory = seen[seen.length - 1].history;
    const lastUserEntry = [...modelHistory].reverse().find((e) => e.role === "user");
    expect(lastUserEntry && lastUserEntry.role === "user" ? lastUserEntry.text : "").toContain(
      'Action log1 "Add a fixed cost" was applied.',
    );
  });
});