freezmo/tests/chat-test.mjs

/**
 * Chat mode (phase 1): the streaming API (validation, allowlist, context
 * travel, history window, thinking separation, usage-limit mapping), the
 * attachment extraction route, and the phone/desktop UI (entry points, scope,
 * streaming, stop, model picker, context picker, edit/regenerate, markdown,
 * math, sources, offline, reduced motion).
 *
 * Runs against a stub Ollama, so no real model is needed.
 */
import http from "node:http";

import { DECK, launchBrowser, makeChecker, makePage, portInUse, seed, spawnApp, waitForHttp } from "./lib.mjs";

const { check, finish } = makeChecker();
const browser = await launchBrowser();

const STUB_PORT = 11997;
const APP_PORT = 3212;
const APP = `http://127.0.0.1:${APP_PORT}`;

/* ------------------------------------------------------------------ */
/* Stub Ollama                                                         */
/* ------------------------------------------------------------------ */

let mode = "ok"; // ok | thinking | slow | limit
let lastChat = null;

const ANSWER = [
  "**Rayleigh** scattering makes the sky blue.",
  "",
  "- short wavelengths scatter more",
  "- violet is mostly absorbed",
  "",
  "| Layer | Effect |",
  "| --- | --- |",
  "| Troposphere | blue |",
  "",
  "```js",
  'const sky = "blue";',
  "```",
  "",
  "Mass-energy: $E = mc^2$ done.",
].join("\n");

const stub = http.createServer((request, response) => {
  if (request.url?.startsWith("/api/tags")) {
    response.writeHead(200, { "Content-Type": "application/json" });
    response.end(JSON.stringify({ models: [{ name: "stub-text" }, { name: "stub-vision" }] }));
    return;
  }
  let raw = "";
  request.on("data", (chunk) => (raw += chunk));
  request.on("end", () => {
    let body = {};
    try {
      body = JSON.parse(raw || "{}");
    } catch {
      /* ignore */
    }
    if (request.url?.startsWith("/api/chat")) {
      const system = String(body.messages?.[0]?.content ?? "");
      // Vision/extraction step (JSON mode).
      if (body.format === "json" || system.includes("transcriber")) {
        response.writeHead(200, { "Content-Type": "application/json" });
        response.end(
          JSON.stringify({
            message: {
              content: JSON.stringify({
                title: "Photo notes",
                subject: "Biology",
                language: "English",
                notes: "EXTRACTED-PHOTO-MARKER: photosynthesis needs light and water.",
                keyConcepts: ["Photosynthesis"],
              }),
            },
          }),
        );
        return;
      }

      lastChat = body;
      if (mode === "limit") {
        response.writeHead(429, { "Content-Type": "application/json" });
        response.end(JSON.stringify({ error: "You have reached your monthly usage limit. Upgrade your plan." }));
        return;
      }

      response.writeHead(200, { "Content-Type": "application/x-ndjson" });
      const chunks = mode === "slow" ? ["Slow ", "answer ", "still ", "coming"] : [ANSWER.slice(0, 24), ANSWER.slice(24, 80), ANSWER.slice(80)];
      let index = 0;
      const writeChunk = () => {
        response.write(JSON.stringify({ message: { content: chunks[index] }, done: false }) + "\n");
        index += 1;
      };
      // Slow mode: the first chunk lands at once so a stop test has content.
      if (mode === "slow") writeChunk();
      const timer = setInterval(
        () => {
          if (index < chunks.length) {
            if (index === 0 && mode === "thinking") {
              response.write(JSON.stringify({ message: { thinking: "REASONING-MARKER let me think…" }, done: false }) + "\n");
            }
            writeChunk();
          } else {
            response.write(JSON.stringify({ message: { content: "" }, done: true }) + "\n");
            response.end();
            clearInterval(timer);
          }
        },
        mode === "slow" ? 700 : 30,
      );
      // `res` close fires when the client aborts or the response ends; `req`
      // close would fire immediately after the request body.
      response.on("close", () => clearInterval(timer));
      return;
    }
    response.writeHead(404);
    response.end();
  });
});
await new Promise((resolve) => stub.listen(STUB_PORT, "127.0.0.1", resolve));

check("stub app port is free", !(await portInUse(APP_PORT)));
const app = spawnApp(APP_PORT, {
  OLLAMA_BASE_URL: `http://127.0.0.1:${STUB_PORT}`,
  OLLAMA_AUTH: "local",
  OLLAMA_VERIFY_KEY: "false",
  CHAT_MODELS: "stub-text,stub-vision",
});

const DECK_NOTES = {
  ...DECK,
  sourceText: "FOTOSYNTEZA-MARKER: plants convert light, water and CO2 into sugar. Chlorophyll absorbs red and blue light.",
};

async function chatPost(payload) {
  const response = await fetch(`${APP}/api/chat`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(payload),
  });
  const events = response.headers.get("content-type")?.includes("ndjson")
    ? (await response.text())
        .trim()
        .split("\n")
        .filter(Boolean)
        .map((line) => JSON.parse(line))
    : null;
  const body = events ? null : await response.json().catch(() => null);
  return { status: response.status, events, body };
}

try {
  const up = await waitForHttp(`${APP}/api/health`);
  check("stub app instance is up", up);

  if (up) {
    /* 1. API: validation + allowlist ---------------------------------- */
    const badModel = await chatPost({ model: "evil-model", messages: [{ role: "user", content: "hi" }] });
    check("api: model outside the allowlist → 400 bad_model", badModel.status === 400 && badModel.body?.kind === "bad_model", JSON.stringify(badModel.body));

    const empty = await chatPost({ messages: [] });
    check("api: empty messages → 400", empty.status === 400);

    const wrongLast = await chatPost({ messages: [{ role: "user", content: "hi" }, { role: "assistant", content: "hello" }] });
    check("api: last message must be from the user → 400", wrongLast.status === 400);

    /* 2. API: streaming + context travel + thinking -------------------- */
    const ok = await chatPost({
      messages: [{ role: "user", content: "Why is the sky blue?" }],
      context: [{ deckId: "seed_deck", deckTitle: "Sky deck", section: "Notes", text: "RAYLEIGH-CONTEXT-MARKER scattering notes." }],
    });
    const tokens = (ok.events ?? []).filter((event) => event.type === "token").map((event) => event.text).join("");
    const done = (ok.events ?? []).find((event) => event.type === "done");
    check("api: streams tokens and a done event", ok.status === 200 && tokens.includes("Rayleigh") && done?.model === "stub-text", done?.model);
    check("api: thinking never mixes into the answer", !tokens.includes("REASONING-MARKER"));
    check(
      "api: system prompt carries the chosen context + the rules",
      lastChat?.messages?.[0]?.content?.includes("RAYLEIGH-CONTEXT-MARKER") &&
        lastChat?.messages?.[0]?.content?.includes("primary source"),
    );
    check("api: the conversation window follows the system prompt", lastChat?.messages?.[1]?.role === "user");

    mode = "thinking";
    const thinking = await chatPost({ messages: [{ role: "user", content: "Think first" }] });
    const thinkingEvents = (thinking.events ?? []).filter((event) => event.type === "thinking").map((event) => event.text).join("");
    const thinkingTokens = (thinking.events ?? []).filter((event) => event.type === "token").map((event) => event.text).join("");
    check("api: reasoning streams separately", thinkingEvents.includes("REASONING-MARKER") && !thinkingTokens.includes("REASONING-MARKER"));
    mode = "ok";

    /* 3. API: history window + context cap ----------------------------- */
    const longHistory = Array.from({ length: 60 }, (_, index) => ({
      role: index % 2 === 0 ? "user" : "assistant",
      content: `Message ${index}: ${"x".repeat(600)}`,
    }));
    longHistory[59] = { role: "user", content: `Last question ${"y".repeat(600)}` };
    await chatPost({ messages: longHistory });
    const sentCount = lastChat?.messages?.length ?? 0;
    check("api: sliding window trims the history", sentCount > 1 && sentCount < 61, `${sentCount} messages reached the model`);

    const manyChunks = Array.from({ length: 40 }, (_, index) => ({
      deckId: `deck_${index}`,
      deckTitle: `Deck ${index}`,
      section: "Notes",
      text: `CHUNK-${index} ${"z".repeat(11_500)}`,
    }));
    await chatPost({ messages: [{ role: "user", content: "question" }], context: manyChunks });
    const systemLength = String(lastChat?.messages?.[0]?.content ?? "").length;
    check("api: context is capped well below what was sent", systemLength > 1000 && systemLength < 60_000, `${systemLength} chars`);
    check("api: the first (best) chunk survives the cap", String(lastChat?.messages?.[0]?.content ?? "").includes("CHUNK-0"));

    /* 4. API: general chat + usage limit ------------------------------- */
    const general = await chatPost({ messages: [{ role: "user", content: "hello" }], context: [] });
    check("api: no context → general-chat system prompt", String(lastChat?.messages?.[0]?.content ?? "").includes("No notes are attached") && general.status === 200);

    mode = "limit";
    const limited = await chatPost({ messages: [{ role: "user", content: "hi" }] });
    const limitEvent = (limited.events ?? []).find((event) => event.type === "error");
    check("api: monthly usage 429 → usage_limit error event", limitEvent?.kind === "usage_limit", JSON.stringify(limitEvent).slice(0, 90));
    mode = "ok";

    /* 5. API: attachment extraction ------------------------------------ */
    const extracted = await fetch(`${APP}/api/chat/extract`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ images: ["data:image/png;base64," + "A".repeat(64)] }),
    });
    const extractedBody = await extracted.json();
    check("extract: photo → vision text", extracted.status === 200 && extractedBody.text?.includes("EXTRACTED-PHOTO-MARKER"));

    const plain = await fetch(`${APP}/api/chat/extract`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ text: "pasted text stays as is" }),
    });
    const plainBody = await plain.json();
    check("extract: plain text needs no model call", plain.status === 200 && plainBody.text === "pasted text stays as is");

    const emptyExtract = await fetch(`${APP}/api/chat/extract`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({}),
    });
    check("extract: nothing attached → 400", emptyExtract.status === 400);

    /* 6. UI: entry points ---------------------------------------------- */
    const desktop = await makePage(browser, { width: 1280, height: 900, deviceScaleFactor: 1 });
    await seed(desktop.page, { decks: [DECK_NOTES] });
    await desktop.page.goto(APP, { waitUntil: "networkidle" });
    await desktop.page.waitForTimeout(900);
    const askLink = desktop.page.getByRole("link", { name: "Ask", exact: true });
    check("desktop: Library has an Ask link", (await askLink.count()) === 1 && (await askLink.getAttribute("href")) === "/chat");
    check("desktop: sidebar has a Chat row", (await desktop.page.locator('aside a[href="/chat"]').count()) === 1);
    await askLink.click();
    await desktop.page.waitForTimeout(1200);
    check("desktop: chat route opens with the sidebar", desktop.page.url().endsWith("/chat") && (await desktop.page.locator("aside").count()) >= 1);
    check("desktop: no tab bar on chat", (await desktop.page.locator("nav.tabbar").count()) === 0);
    await desktop.context.close();

    const phone = await makePage(browser);
    await seed(phone.page, { decks: [DECK_NOTES] });
    await phone.page.goto(APP, { waitUntil: "networkidle" });
    await phone.page.waitForTimeout(900);
    const askButton = phone.page.getByRole("link", { name: "Ask", exact: true });
    check("phone: Library has an Ask button", (await askButton.count()) === 1 && (await askButton.getAttribute("href")) === "/chat");
    await phone.page.getByRole("button", { name: "Deck actions" }).click();
    await phone.page.waitForTimeout(300);
    const askDeck = phone.page.getByRole("link", { name: "Ask about this deck" });
    check("phone: deck menu has Ask about this deck", (await askDeck.count()) === 1 && (await askDeck.getAttribute("href"))?.includes("/chat?deck=seed_deck"));
    await askDeck.click();
    await phone.page.waitForTimeout(1400);
    check("phone: chat opens scoped to the deck", phone.page.url().includes("deck=seed_deck") && (await phone.page.getByText("1 deck").count()) >= 1);
    check("phone: no tab bar on chat", (await phone.page.locator("nav.tabbar").count()) === 0);
    check("phone: back button present", (await phone.page.getByRole("button", { name: "Back" }).count()) === 1);

    /* 7. UI: streaming + markdown + math + sources ---------------------- */
    check("chat: empty state with suggestions", await phone.page.getByText("Ask your notes anything").isVisible());
    check("chat: deck suggestion present", (await phone.page.getByRole("button", { name: /Summarize my notes/ }).count()) === 1);
    const context = phone.page.locator(".chat-md").first();
    await phone.page.getByRole("button", { name: /Summarize my notes/ }).click();
    await phone.page.waitForTimeout(2500);
    check("chat: user bubble appears", (await phone.page.getByText("Summarize my notes").count()) >= 1);
    check("chat: assistant answer streams in", await context.getByText("scattering makes the sky blue").isVisible());
    check("chat: markdown bold + list render", (await context.locator("strong").count()) >= 1 && (await context.locator("li").count()) >= 2);
    check("chat: table renders", (await context.locator("table").count()) === 1);
    check(
      "chat: code block + copy button",
      (await context.locator("pre").count()) === 1 && (await context.getByRole("button", { name: "Copy" }).count()) >= 1,
    );
    check("chat: math renders through KaTeX", (await phone.page.locator(".katex").count()) >= 1);
    check("chat: sources chip links to the deck", (await phone.page.locator('a[href="/deck/seed_deck"]').count()) >= 1);
    check("chat: context reached the model", String(lastChat?.messages?.[0]?.content ?? "").includes("FOTOSYNTEZA-MARKER"));

    /* 8. UI: model picker + context picker ------------------------------ */
    await phone.page.getByRole("button", { name: "stub-text" }).click();
    await phone.page.waitForTimeout(400);
    check("model: picker lists both models with availability", (await phone.page.getByText("stub-vision").count()) >= 1 && (await phone.page.getByText("available").count()) >= 1);
    await phone.page.getByRole("button", { name: /stub-vision/ }).click();
    await phone.page.waitForTimeout(400);
    check("model: choice updates the chip", (await phone.page.getByRole("button", { name: "stub-vision" }).count()) === 1);
    check(
      "model: choice is remembered",
      (await phone.page.evaluate(() => localStorage.getItem("freezmo.chat.model.v1"))) === "stub-vision",
    );
    await phone.page.reload({ waitUntil: "networkidle" });
    await phone.page.waitForTimeout(1200);
    check("model: remembered across reloads", (await phone.page.getByRole("button", { name: "stub-vision" }).count()) === 1);

    await phone.page.getByRole("button", { name: "1 deck" }).click();
    await phone.page.waitForTimeout(400);
    check("scope: picker opens with the modes", (await phone.page.getByText("All notes").count()) >= 1 && (await phone.page.getByText("No notes").count()) >= 1);
    await phone.page.getByRole("button", { name: /No notes/ }).first().click();
    await phone.page.waitForTimeout(400);
    check("scope: chip follows the choice", (await phone.page.getByRole("button", { name: "No notes" }).count()) === 1);

    /* 9. UI: thinking row ----------------------------------------------- */
    mode = "thinking";
    await phone.page.getByRole("textbox").fill("Think about it");
    await phone.page.getByRole("button", { name: "Send", exact: true }).click();
    await phone.page.waitForTimeout(2500);
    const thinkingRow = phone.page.getByRole("button", { name: /Thinking/ }).last();
    check("thinking: row appears", (await thinkingRow.count()) >= 1);
    check("thinking: reasoning hidden by default", (await phone.page.getByText("REASONING-MARKER").count()) === 0);
    await thinkingRow.click();
    await phone.page.waitForTimeout(300);
    check("thinking: expanding reveals it", (await phone.page.getByText(/REASONING-MARKER/).count()) >= 1);
    mode = "ok";

    /* 10. UI: stop mid-stream ------------------------------------------- */
    mode = "slow";
    await phone.page.getByRole("textbox").fill("Slow one");
    await phone.page.getByRole("button", { name: "Send", exact: true }).click();
    await phone.page.waitForTimeout(1700);
    check("stop: stop button appears while streaming", (await phone.page.getByRole("button", { name: "Stop" }).count()) === 1);
    await phone.page.getByRole("button", { name: "Stop" }).click();
    await phone.page.waitForTimeout(700);
    check("stop: send button returns", (await phone.page.getByRole("button", { name: "Send", exact: true }).count()) === 1);
    const partial = await phone.page.locator(".chat-md").last().innerText().catch(() => "");
    // Next's route announcer is also role=alert — only real error cards count.
    const alerts = await phone.page.locator('[role="alert"]:not(#__next-route-announcer__)').count();
    check("stop: partial answer kept, no error card", partial.includes("Slow") && alerts === 0, JSON.stringify(partial.slice(0, 30)));
    mode = "ok";

    /* 11. UI: regenerate + edit ----------------------------------------- */
    const lastAssistantActions = phone.page.locator("div").filter({ has: phone.page.getByRole("button", { name: "Regenerate" }) });
    check("actions: regenerate offered on the last answer", (await phone.page.getByRole("button", { name: "Regenerate" }).count()) >= 1);
    const before = await phone.page.locator(".chat-md").count();
    await phone.page.getByRole("button", { name: "Regenerate" }).last().click();
    await phone.page.waitForTimeout(2200);
    check("actions: regenerate streams a fresh answer", (await phone.page.locator(".chat-md").count()) === before);
    void lastAssistantActions;
    check("actions: edit offered on the last user message", (await phone.page.getByRole("button", { name: "Edit and resend" }).count()) === 1);
    await phone.page.getByRole("button", { name: "Edit and resend" }).click();
    await phone.page.waitForTimeout(400);
    check("actions: edit restores the text into the composer", (await phone.page.getByRole("textbox").inputValue()) === "Slow one");

    /* 12. UI: usage limit + offline ------------------------------------- */
    mode = "limit";
    await phone.page.getByRole("textbox").fill("Will hit the limit");
    await phone.page.getByRole("button", { name: "Send", exact: true }).click();
    await phone.page.waitForTimeout(2000);
    check("limit: error card shown", (await phone.page.getByRole("alert").count()) >= 1 && (await phone.page.getByText(/free monthly Ollama usage/i).count()) >= 1);
    check("limit: Settings + Ollama actions offered", (await phone.page.getByRole("button", { name: /Open settings/ }).count()) >= 1 && (await phone.page.getByRole("button", { name: /Ollama/ }).count()) >= 1);
    check("limit: the draft is kept for a retry", (await phone.page.getByRole("textbox").inputValue()) === "Will hit the limit");
    mode = "ok";

    await phone.context.setOffline(true);
    await phone.page.waitForTimeout(500);
    check("offline: send disabled with a reason", (await phone.page.getByRole("button", { name: "Send", exact: true }).isDisabled()) && (await phone.page.getByText(/Offline — sending is paused/).count()) === 1);
    await phone.context.setOffline(false);
    await phone.page.waitForTimeout(400);

    /* 13. UI: attachment ------------------------------------------------- */
    await phone.page.getByRole("button", { name: "Attach photos or documents" }).click();
    await phone.page.locator('input[type="file"]').setInputFiles({
      name: "notes.png",
      mimeType: "image/png",
      buffer: Buffer.from(
        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==",
        "base64",
      ),
    });
    await phone.page.waitForTimeout(1800);
    check("attach: chip appears after extraction", (await phone.page.getByText("notes.png").count()) >= 1);
    await phone.page.getByRole("textbox").fill("What does the photo say?");
    await phone.page.getByRole("button", { name: "Send", exact: true }).click();
    await phone.page.waitForTimeout(2200);
    check("attach: extracted text travels as context", String(lastChat?.messages?.[0]?.content ?? "").includes("EXTRACTED-PHOTO-MARKER"));

    /* 14. UI: new chat + old-decks note --------------------------------- */
    await phone.page.getByRole("button", { name: "New chat" }).click();
    await phone.page.waitForTimeout(500);
    check("new chat: clears the conversation", (await phone.page.locator(".chat-md").count()) === 0 && (await phone.page.getByText("Ask your notes anything").count()) === 1);

    const oldDeck = { ...DECK, id: "old_deck", title: "Imported deck", sourceText: "" };
    const oldPage = await makePage(browser);
    await seed(oldPage.page, { decks: [oldDeck] });
    await oldPage.page.goto(`${APP}/chat?deck=old_deck`, { waitUntil: "networkidle" });
    await oldPage.page.waitForTimeout(1200);
    check("old decks: the fallback notice is shown", (await oldPage.page.getByText(/Older decks only have the generated study content/).count()) === 1);
    check("old decks: scope chip scoped to it", (await oldPage.page.getByText("1 deck").count()) >= 1);
    await oldPage.context.close();

    /* 15. reduced motion still works ------------------------------------ */
    const reduced = await makePage(browser, { reducedMotion: "reduce" });
    await seed(reduced.page, { decks: [DECK_NOTES] });
    await reduced.page.goto(`${APP}/chat`, { waitUntil: "networkidle" });
    await reduced.page.waitForTimeout(1000);
    await reduced.page.getByRole("textbox").fill("Reduced motion question");
    await reduced.page.getByRole("button", { name: "Send", exact: true }).click();
    await reduced.page.waitForTimeout(2500);
    check("reduced motion: chat still answers", (await reduced.page.locator(".chat-md").count()) >= 1);
    await reduced.context.close();

    await phone.context.close();
  }
} finally {
  app.stop();
  stub.close();
  await browser.close();
}

finish();