PubScript/packages/ai/src/ai-session.test.ts

import { afterEach, beforeEach, describe, expect, it } from "vitest";

import { AIStore } from "./index.js";
import type { SessionRow } from "./index.js";
import type { AISessionPersistence } from "./index.js";

/** A simple in-memory persistence backend for tests. */
function memoryPersistence(): AISessionPersistence {
  const rows = new Map<string, SessionRow>();
  return {
    list: async (projectId) =>
      Array.from(rows.values())
        .filter((r) => r.projectId === projectId)
        .map((r) => ({ ...r })),
    get: async (id) => (rows.has(id) ? ({ ...rows.get(id) } as SessionRow) : null),
    save: async (session) => {
      rows.set(session.id, { ...session });
    },
    delete: async (id) => {
      rows.delete(id);
    },
  };
}

describe("AIStore", () => {
  let store: AIStore;
  let persistence: AISessionPersistence;

  beforeEach(() => {
    persistence = memoryPersistence();
    store = new AIStore("proj-1", persistence);
  });

  afterEach(async () => {
    await store.dispose();
  });

  it("creates a session with a generated id and timestamp", async () => {
    const session = await store.createSession("Chapter 1");
    expect(session.id).toBeTruthy();
    expect(session.title).toBe("Chapter 1");
    expect(session.projectId).toBe("proj-1");
    expect(session.createdAt).toBeTruthy();
    expect(session.updatedAt).toBe(session.createdAt);
    expect(session.messages).toEqual([]);
    expect(session.suggestions).toEqual([]);
    expect(session.notes).toEqual([]);
  });

  it("auto-generates a title when none is provided", async () => {
    const session = await store.createSession();
    expect(session.title).toMatch(/^New AI session \d+$/);
  });

  it("lists sessions in reverse chronological order", async () => {
    const a = await store.createSession("A");
    const b = await store.createSession("B");
    const c = await store.createSession("C");
    const list = await store.listSessions();
    expect(list.map((s) => s.id)).toEqual([c.id, b.id, a.id]);
  });

  it("adds a message to a session", async () => {
    const session = await store.createSession();
    const updated = await store.addMessage(session.id, {
      role: "user",
      content: "hello",
    });
    expect(updated.messages).toHaveLength(1);
    expect(updated.messages[0]?.role).toBe("user");
  });

  it("adds a suggestion and lists it", async () => {
    const session = await store.createSession();
    const updated = await store.addSuggestion(session.id, {
      sessionId: session.id,
      kind: "brainstorm",
      title: "Ideas",
      text: "some ideas",
      summary: "summary",
      status: "pending",
    });
    expect(updated.suggestions).toHaveLength(1);
    const list = await store.listSuggestions(session.id);
    expect(list).toHaveLength(1);
    expect(list[0]?.kind).toBe("brainstorm");
  });

  it("applies a suggestion (marks applied, updates updatedAt)", async () => {
    const session = await store.createSession();
    const withSuggestion = await store.addSuggestion(session.id, {
      sessionId: session.id,
      kind: "brainstorm",
      title: "Ideas",
      text: "some ideas",
      summary: "summary",
      status: "pending",
    });
    const suggestionId = withSuggestion.suggestions[0]?.id;
    const updated = await store.applySuggestion(session.id, suggestionId!);
    expect(updated.suggestions[0]?.status).toBe("applied");
    expect(updated.suggestions[0]?.appliedAt).toBeTruthy();
  });

  it("resolves a suggestion (marks resolved)", async () => {
    const session = await store.createSession();
    const withSuggestion = await store.addSuggestion(session.id, {
      sessionId: session.id,
      kind: "brainstorm",
      title: "Ideas",
      text: "some ideas",
      summary: "summary",
      status: "pending",
    });
    const suggestionId = withSuggestion.suggestions[0]?.id;
    const updated = await store.resolveSuggestion(session.id, suggestionId!, true);
    expect(updated.suggestions[0]?.status).toBe("accepted");
    expect(updated.suggestions[0]?.resolvedAt).toBeTruthy();
  });

  it("adds a note and updates it", async () => {
    const session = await store.createSession();
    const withNote = await store.addNote(session.id, {
      sessionId: session.id,
      kind: "character",
      title: "Aria",
      body: "the protagonist",
      tags: ["main"],
      status: "draft",
    });
    const noteId = withNote.notes[0]?.id;
    expect(withNote.notes).toHaveLength(1);
    const updated = await store.updateNote(session.id, noteId!, { body: "the protagonist, a mage" });
    expect(updated.notes[0]?.body).toBe("the protagonist, a mage");
    expect(updated.notes[0]?.tags).toEqual(["main"]);
  });

  it("merges notes into a single note", async () => {
    const session = await store.createSession();
    const note1 = await store.addNote(session.id, {
      sessionId: session.id,
      kind: "character",
      title: "Aria",
      body: "part one",
      tags: ["a"],
      status: "draft",
    });
    const note2 = await store.addNote(session.id, {
      sessionId: session.id,
      kind: "character",
      title: "Notes",
      body: "part two",
      tags: ["b"],
      status: "draft",
    });
    const merged = await store.mergeNotes(session.id, [note1.notes[0]!.id, note2.notes[1]!.id]);
    expect(merged.notes).toHaveLength(1);
    expect(merged.notes[0]?.body).toBe("part one\n\npart two");
    expect(merged.notes[0]?.tags).toEqual(["a", "b"]);
  });

  it("deletes a session", async () => {
    const session = await store.createSession();
    await store.deleteSession(session.id);
    const list = await store.listSessions();
    expect(list).toHaveLength(0);
  });

  it("throws when mutating a missing session", async () => {
    await expect(store.getSession("missing")).resolves.toBeNull();
    await expect(store.addMessage("missing", {
      role: "user", content: "x",
    })).rejects.toThrow();
  });

  it("persists across a fresh store instance sharing the same backend", async () => {
    const session = await store.createSession("persisted");
    const other = new AIStore("proj-1", persistence);
    const list = await other.listSessions();
    expect(list.map((s) => s.id)).toContain(session.id);
    await other.dispose();
  });
});