evOo-local/packages/db/src/studio-agent.postgres.test.ts

import { randomUUID } from "node:crypto";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { bootstrapUserSpace } from "./bootstrap-user.js";
import { createDb, type PrismaClient } from "./client.js";
import { ensureMainBot } from "./main-agent.js";
import {
  ensureStudioBot,
  STUDIO_BOT_DEFAULTS,
  STUDIO_BOT_TOOLS,
  studioBotSpawnKey,
} from "./studio-agent.js";

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

const signupEnv = { signupsEnabled: "true", signupAllowlist: "" };

describePostgres("Studio specialist (PostgreSQL)", () => {
  let prisma: PrismaClient;
  let close: () => Promise<void>;
  const userIds: string[] = [];

  async function newUser(name = "Ada") {
    const id = `st${randomUUID().replaceAll("-", "")}`;
    userIds.push(id);
    await prisma.user.create({ data: { id, name, email: `${id}@rakazo.test` } });
    const { spaceId } = await bootstrapUserSpace(prisma, { id }, signupEnv, {
      claimDeploymentOwner: false,
    });
    return { userId: id, homeSpaceId: spaceId };
  }

  beforeAll(() => {
    const db = createDb(databaseUrl!);
    prisma = db.prisma;
    close = async () => {
      await db.prisma.$disconnect();
      await db.pool.end();
    };
  });

  afterAll(async () => {
    if (!prisma) return;
    const memberships = await prisma.member.findMany({
      where: { userId: { in: userIds } },
      select: { organizationId: true },
    });
    await prisma.organization.deleteMany({
      where: { id: { in: memberships.map((member) => member.organizationId) } },
    });
    await prisma.user.deleteMany({ where: { id: { in: userIds } } });
    await close();
  });

  it("creates one Studio next to the main bot, even under concurrent calls", async () => {
    const { userId, homeSpaceId } = await newUser();
    const results = await Promise.all([
      ensureStudioBot(prisma, userId),
      ensureStudioBot(prisma, userId),
      ensureStudioBot(prisma, userId),
    ]);
    expect(new Set(results.map((bot) => bot?.id)).size).toBe(1);
    const main = await ensureMainBot(prisma, userId);
    const studios = await prisma.bot.findMany({ where: { userId, kind: "specialist" } });
    expect(studios).toHaveLength(1);
    expect(studios[0]).toMatchObject({
      id: results[0]?.id,
      spaceId: homeSpaceId,
      name: "Studio",
      instructions: STUDIO_BOT_DEFAULTS.instructions,
      tags: ["media"],
      restrictBuiltinTools: true,
      allowedBuiltinTools: [...STUDIO_BOT_TOOLS],
      spawnKey: studioBotSpawnKey(userId),
    });
    expect(main.spaceId).toBe(homeSpaceId);
    // A thread makes it addressable: the main bot's specialist directory lists bots with one.
    expect(await prisma.thread.count({ where: { botId: studios[0]?.id } })).toBe(1);
    expect(await ensureStudioBot(prisma, userId)).toEqual(results[0]);
  });

  it("gives every member of a shared space their own Studio", async () => {
    const owner = await newUser();
    const ownerStudio = await ensureStudioBot(prisma, owner.userId);
    const guestId = `st${randomUUID().replaceAll("-", "")}`;
    userIds.push(guestId);
    await prisma.user.create({
      data: { id: guestId, name: "Guest", email: `${guestId}@rakazo.test` },
    });
    const organization = await prisma.space.findUniqueOrThrow({
      where: { id: owner.homeSpaceId },
      select: { organizationId: true },
    });
    // The organization's default-space trigger makes the guest a member of the owner's space.
    await prisma.member.create({
      data: {
        id: `member-${guestId}`,
        organizationId: organization.organizationId,
        userId: guestId,
        role: "member",
        createdAt: new Date(),
      },
    });
    const guestStudio = await ensureStudioBot(prisma, guestId);
    expect(guestStudio).not.toBeNull();
    expect(guestStudio?.id).not.toBe(ownerStudio?.id);
    expect(await prisma.bot.findUniqueOrThrow({ where: { id: guestStudio!.id } })).toMatchObject({
      userId: guestId,
      spawnKey: studioBotSpawnKey(guestId),
    });
  });

  it("does not bring back a Studio the user archived or deleted", async () => {
    const archiver = await newUser();
    const archived = await ensureStudioBot(prisma, archiver.userId);
    await prisma.bot.update({ where: { id: archived!.id }, data: { archivedAt: new Date() } });
    await expect(ensureStudioBot(prisma, archiver.userId)).resolves.toBeNull();
    expect(await prisma.bot.count({ where: { userId: archiver.userId, name: "Studio" } })).toBe(1);

    const deleter = await newUser();
    const deleted = await ensureStudioBot(prisma, deleter.userId);
    const row = await prisma.bot.findUniqueOrThrow({ where: { id: deleted!.id } });
    // What destroyBot records when the user deletes a bot.
    await prisma.botDeletion.create({
      data: {
        id: row.id,
        spaceId: row.spaceId,
        name: row.name,
        spawnKey: row.spawnKey,
        deletedByUserId: deleter.userId,
        memoriesPreserved: true,
      },
    });
    await prisma.bot.delete({ where: { id: row.id } });
    await expect(ensureStudioBot(prisma, deleter.userId)).resolves.toBeNull();
    expect(await prisma.bot.count({ where: { userId: deleter.userId, name: "Studio" } })).toBe(0);
  });
});