stack/apps/api/src/agent-memory/agent-memory.integration.spec.ts

import { beforeAll, beforeEach, describe, expect, it, afterAll } from "vitest";
import { randomUUID as uuid } from "crypto";
import { Test, TestingModule } from "@nestjs/testing";
import { NotFoundException } from "@nestjs/common";
import { PrismaClient } from "@prisma/client";
import { AgentMemoryService } from "./agent-memory.service";
import { PrismaService } from "../prisma/prisma.service";

const shouldRunDbIntegrationTests =
  process.env.RUN_DB_TESTS === "true" && Boolean(process.env.DATABASE_URL);
const describeFn = shouldRunDbIntegrationTests ? describe : describe.skip;

async function createWorkspace(
  prisma: PrismaClient,
  label: string
): Promise<{ workspaceId: string; ownerId: string }> {
  const workspace = await prisma.workspace.create({
    data: {
      name: `${label} ${Date.now()}`,
      owner: {
        create: {
          email: `${label.toLowerCase().replace(/\s+/g, "-")}-${Date.now()}@example.com`,
          name: `${label} Owner`,
        },
      },
    },
  });

  return {
    workspaceId: workspace.id,
    ownerId: workspace.ownerId,
  };
}

describeFn("AgentMemoryService Integration", () => {
  let moduleRef: TestingModule;
  let prisma: PrismaClient;
  let service: AgentMemoryService;
  let setupComplete = false;

  let workspaceAId: string;
  let workspaceAOwnerId: string;
  let workspaceBId: string;
  let workspaceBOwnerId: string;

  beforeAll(async () => {
    prisma = new PrismaClient();
    await prisma.$connect();

    const workspaceA = await createWorkspace(prisma, "Agent Memory Integration A");
    workspaceAId = workspaceA.workspaceId;
    workspaceAOwnerId = workspaceA.ownerId;

    const workspaceB = await createWorkspace(prisma, "Agent Memory Integration B");
    workspaceBId = workspaceB.workspaceId;
    workspaceBOwnerId = workspaceB.ownerId;

    moduleRef = await Test.createTestingModule({
      providers: [
        AgentMemoryService,
        {
          provide: PrismaService,
          useValue: prisma,
        },
      ],
    }).compile();

    service = moduleRef.get<AgentMemoryService>(AgentMemoryService);
    setupComplete = true;
  });

  beforeEach(async () => {
    if (!setupComplete) {
      return;
    }

    await prisma.agentMemory.deleteMany({
      where: {
        workspaceId: {
          in: [workspaceAId, workspaceBId],
        },
      },
    });
  });

  afterAll(async () => {
    if (!prisma) {
      return;
    }

    const workspaceIds = [workspaceAId, workspaceBId].filter(
      (id): id is string => typeof id === "string"
    );
    const ownerIds = [workspaceAOwnerId, workspaceBOwnerId].filter(
      (id): id is string => typeof id === "string"
    );

    if (workspaceIds.length > 0) {
      await prisma.agentMemory.deleteMany({
        where: {
          workspaceId: {
            in: workspaceIds,
          },
        },
      });
      await prisma.workspace.deleteMany({ where: { id: { in: workspaceIds } } });
    }

    if (ownerIds.length > 0) {
      await prisma.user.deleteMany({ where: { id: { in: ownerIds } } });
    }

    if (moduleRef) {
      await moduleRef.close();
    }
    await prisma.$disconnect();
  });

  it("upserts and lists memory entries", async () => {
    if (!setupComplete) {
      return;
    }

    const agentId = `agent-${uuid()}`;

    const entry = await service.upsert(workspaceAId, agentId, "session-context", {
      value: { intent: "create-tests", depth: "integration" },
    });

    expect(entry.workspaceId).toBe(workspaceAId);
    expect(entry.agentId).toBe(agentId);
    expect(entry.key).toBe("session-context");

    const listed = await service.findAll(workspaceAId, agentId);

    expect(listed).toHaveLength(1);
    expect(listed[0]?.id).toBe(entry.id);
    expect(listed[0]?.value).toMatchObject({ intent: "create-tests" });
  });

  it("updates existing key via upsert without creating duplicates", async () => {
    if (!setupComplete) {
      return;
    }

    const agentId = `agent-${uuid()}`;

    const first = await service.upsert(workspaceAId, agentId, "preferences", {
      value: { model: "fast" },
    });

    const second = await service.upsert(workspaceAId, agentId, "preferences", {
      value: { model: "accurate" },
    });

    expect(second.id).toBe(first.id);
    expect(second.value).toMatchObject({ model: "accurate" });

    const rowCount = await prisma.agentMemory.count({
      where: {
        workspaceId: workspaceAId,
        agentId,
        key: "preferences",
      },
    });

    expect(rowCount).toBe(1);
  });

  it("lists keys in sorted order and isolates by workspace", async () => {
    if (!setupComplete) {
      return;
    }

    const agentId = `agent-${uuid()}`;

    await service.upsert(workspaceAId, agentId, "beta", { value: { v: 2 } });
    await service.upsert(workspaceAId, agentId, "alpha", { value: { v: 1 } });
    await service.upsert(workspaceBId, agentId, "alpha", { value: { v: 99 } });

    const workspaceAEntries = await service.findAll(workspaceAId, agentId);
    const workspaceBEntries = await service.findAll(workspaceBId, agentId);

    expect(workspaceAEntries.map((row) => row.key)).toEqual(["alpha", "beta"]);
    expect(workspaceBEntries).toHaveLength(1);
    expect(workspaceBEntries[0]?.value).toMatchObject({ v: 99 });
  });

  it("throws NotFoundException when requesting unknown key", async () => {
    if (!setupComplete) {
      return;
    }

    await expect(service.findOne(workspaceAId, `agent-${uuid()}`, "missing")).rejects.toThrow(
      NotFoundException
    );
  });
});