evOo-local/packages/memory/src/commit.postgres.test.ts

import type { AdapterContext, MemoryCommitRequest } from "@rakazo/adapter-kit";
import { MEMORY_DOCUMENT_MAX_BYTES, MemoryDocumentTooLargeError } from "@rakazo/adapter-kit";
import { createDb, type PrismaClient } from "@rakazo/db";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { MarkdownMemoryStore } from "./index.js";

const databaseUrl = process.env.DATABASE_URL;
const describePostgres =
  process.env.VERIFY_DATABASE && databaseUrl ? describe.sequential : describe.skip;

describePostgres("memory commits (PostgreSQL)", () => {
  const context: AdapterContext = {
    operationId: "memory-commit-test",
    traceId: "memory-commit-test",
    spaceId: "memory-commit-space",
    userId: "memory-commit-user",
    signal: new AbortController().signal,
  };
  let prisma: PrismaClient;
  let store: MarkdownMemoryStore;
  let close: () => Promise<void>;

  beforeAll(async () => {
    const db = createDb(databaseUrl!);
    prisma = db.prisma;
    store = new MarkdownMemoryStore(prisma);
    close = async () => {
      await prisma.$disconnect();
      await db.pool.end();
    };
    await prisma.organization.create({
      data: {
        id: context.spaceId,
        name: "Memory Test Organization",
        slug: context.spaceId,
        createdAt: new Date(),
        spaces: {
          create: {
            id: context.spaceId,
            name: "Memory Test Space",
            bots: {
              create: {
                id: "memory-commit-bot",
                userId: context.userId,
                name: "Memory Test Bot",
                color: "ink",
              },
            },
          },
        },
      },
    });
  });

  afterAll(async () => {
    if (!prisma) return;
    try {
      await prisma.organization.deleteMany({ where: { id: context.spaceId } });
    } finally {
      await close();
    }
  });

  function request(path: string, content: string): MemoryCommitRequest {
    return {
      scope: "user",
      path,
      content,
      sourceRunId: `run-${content}`,
      sourceThreadId: `thread-${content}`,
    };
  }

  async function document(path: string) {
    return prisma.memoryDocument.findFirstOrThrow({
      where: { spaceId: context.spaceId, path },
      include: { revisions: { orderBy: { revision: "asc" } } },
    });
  }

  it.each(["user", "bot"] as const)(
    "retries overlapping %s saves with fresh revisions and matching history",
    async (scope) => {
      const path = `concurrent-${scope}.md`;
      const write = (content: string): MemoryCommitRequest => ({
        ...request(path, content),
        scope,
        botId: scope === "bot" ? "memory-commit-bot" : undefined,
      });
      const initial = await store.commit(write("initial"), context);
      let reads = 0;
      let release!: () => void;
      const bothRead = new Promise<void>((resolve) => {
        release = resolve;
      });
      const concurrentStore = new MarkdownMemoryStore(
        prisma.$extends({
          query: {
            memoryDocument: {
              async findFirst({ args, query }) {
                const result = await query(args);
                reads += 1;
                // Both first attempts see revision 1 before either can write.
                if (reads === 2) release();
                if (reads <= 2) await bothRead;
                return result;
              },
            },
          },
        }) as PrismaClient,
      );

      const results = await Promise.allSettled([
        concurrentStore.commit(write("first"), context),
        concurrentStore.commit(write("second"), context),
      ]);

      expect(results.map((result) => result.status)).toEqual(["fulfilled", "fulfilled"]);
      expect(reads).toBe(3);
      const saved = results.flatMap((result) =>
        result.status === "fulfilled" ? [result.value] : [],
      );
      expect(saved.map((result) => result.revision).sort()).toEqual([2, 3]);
      const current = await document(path);
      expect(current.revision).toBe(3);
      expect(current.content).toBe(saved.find((result) => result.revision === 3)!.content);
      expect(current.revisions).toHaveLength(3);
      for (const result of [initial, ...saved]) {
        expect(
          current.revisions.find((revision) => revision.revision === result.revision),
        ).toMatchObject({
          documentId: result.id,
          revision: result.revision,
          content: result.content,
          sourceRunId: `run-${result.content}`,
          sourceThreadId: `thread-${result.content}`,
        });
      }
    },
  );

  it.each([false, true])(
    "rolls back a rejected history insert and permits a caller retry (existing: %s)",
    async (existing) => {
      const path = `rejected-${existing}.md`;
      if (existing) await store.commit(request(path, "initial"), context);
      const before = await prisma.memoryDocument.findMany({
        where: { spaceId: context.spaceId, path },
        include: { revisions: true },
      });
      let inserts = 0;
      const rejectingStore = new MarkdownMemoryStore(
        prisma.$extends({
          query: {
            memoryRevision: {
              async create({ args, query }) {
                inserts += 1;
                // Make PostgreSQL reject the history row after the document write.
                return query({
                  ...args,
                  data: {
                    ...args.data,
                    document: undefined,
                    documentId: "missing-memory-document",
                  },
                });
              },
            },
          },
        }) as PrismaClient,
      );

      await expect(rejectingStore.commit(request(path, "rejected"), context)).rejects.toMatchObject(
        {
          code: "P2003",
        },
      );
      expect(inserts).toBe(1);
      expect(
        await prisma.memoryDocument.findMany({
          where: { spaceId: context.spaceId, path },
          include: { revisions: true },
        }),
      ).toEqual(before);

      const retried = await store.commit(request(path, "rejected"), context);
      expect(retried.revision).toBe(existing ? 2 : 1);
      const current = await document(path);
      expect(current.revisions).toHaveLength(retried.revision);
      expect(current.revisions.at(-1)).toMatchObject({
        revision: retried.revision,
        content: retried.content,
      });
    },
  );

  it("appends to the current content with a revision per note, creating a missing document", async () => {
    const path = "appended.md";
    const append = (content: string): MemoryCommitRequest => ({
      ...request(path, content),
      mode: "append",
    });
    const first = await store.commit(append("- first"), context);
    expect(first).toMatchObject({ revision: 1, content: "- first\n" });
    await store.commit(append("- second"), context);
    const current = await document(path);
    expect(current).toMatchObject({ revision: 2, content: "- first\n- second\n" });
    expect(current.revisions.map(({ revision, content }) => ({ revision, content }))).toEqual([
      { revision: 1, content: "- first\n" },
      { revision: 2, content: "- first\n- second\n" },
    ]);

    // A replace still rewrites the whole document.
    await store.commit(request(path, "# Condensed\n"), context);
    await store.commit(append("- third"), context);
    expect((await document(path)).content).toBe("# Condensed\n\n- third\n");
  });

  it("keeps both notes when two appends overlap", async () => {
    const path = "concurrent-append.md";
    await store.commit(request(path, "# Notes\n"), context);
    let reads = 0;
    let release!: () => void;
    const bothRead = new Promise<void>((resolve) => {
      release = resolve;
    });
    const concurrentStore = new MarkdownMemoryStore(
      prisma.$extends({
        query: {
          memoryDocument: {
            async findFirst({ args, query }) {
              const result = await query(args);
              reads += 1;
              if (reads === 2) release();
              if (reads <= 2) await bothRead;
              return result;
            },
          },
        },
      }) as PrismaClient,
    );

    const results = await Promise.allSettled([
      concurrentStore.commit({ ...request(path, "- a"), mode: "append" }, context),
      concurrentStore.commit({ ...request(path, "- b"), mode: "append" }, context),
    ]);

    expect(results.map((result) => result.status)).toEqual(["fulfilled", "fulfilled"]);
    const current = await document(path);
    expect(current.revision).toBe(3);
    expect(current.content).toMatch(/^# Notes\n\n- [ab]\n- [ab]\n$/);
    expect(current.content).toContain("- a\n");
    expect(current.content).toContain("- b\n");
  });

  it("refuses a document past the size limit without writing anything", async () => {
    const path = "full.md";
    const nearlyFull = `${"x".repeat(MEMORY_DOCUMENT_MAX_BYTES - 10)}\n`;
    await store.commit(request(path, nearlyFull), context);

    await expect(
      store.commit({ ...request(path, "- one note too many"), mode: "append" }, context),
    ).rejects.toBeInstanceOf(MemoryDocumentTooLargeError);
    await expect(
      store.commit(request(path, "é".repeat(MEMORY_DOCUMENT_MAX_BYTES / 2 + 1)), context),
    ).rejects.toBeInstanceOf(MemoryDocumentTooLargeError);
    const current = await document(path);
    expect(current).toMatchObject({ revision: 1, content: nearlyFull });
    expect(current.revisions).toHaveLength(1);

    await store.commit({ ...request(path, "- ok"), mode: "append" }, context);
    expect((await document(path)).revision).toBe(2);
  });

  it("bounds conflict retries, preserves only committed saves, and allows a later retry", async () => {
    const path = "exhausted.md";
    await store.commit(request(path, "initial"), context);
    let attempts = 0;
    const conflictingStore = new MarkdownMemoryStore(
      prisma.$extends({
        query: {
          memoryDocument: {
            async findFirst({ args, query }) {
              const result = await query(args);
              attempts += 1;
              // Invalidate every attempted snapshot with a separate committed save.
              await store.commit(request(path, `winner-${attempts}`), context);
              return result;
            },
          },
        },
      }) as PrismaClient,
    );

    await expect(conflictingStore.commit(request(path, "rejected"), context)).rejects.toMatchObject(
      {
        code: "P2034",
      },
    );
    expect(attempts).toBe(3);
    const current = await document(path);
    expect(current).toMatchObject({ revision: 4, content: "winner-3" });
    expect(current.revisions.map(({ revision, content }) => ({ revision, content }))).toEqual([
      { revision: 1, content: "initial" },
      { revision: 2, content: "winner-1" },
      { revision: 3, content: "winner-2" },
      { revision: 4, content: "winner-3" },
    ]);

    const retried = await store.commit(request(path, "rejected"), context);
    expect(retried.revision).toBe(5);
    const afterRetry = await document(path);
    expect(afterRetry).toMatchObject({ revision: 5, content: "rejected" });
    expect(afterRetry.revisions).toHaveLength(5);
    expect(afterRetry.revisions.at(-1)).toMatchObject({ revision: 5, content: "rejected" });
  });
});