evOo-local/packages/adapters/src/profile-consolidation.postgres.test.ts

import { randomUUID } from "node:crypto";
import type { AgentRunRequest, AgentRuntime, BackgroundJob } from "@rakazo/adapter-kit";
import type { PrismaClient } from "@rakazo/db";
import { bootstrapUserSpace, createDb, ensureMainBot, USER_FACT_LIMIT } from "@rakazo/db";
import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
import { EmulatorEmbeddingProvider, emulatedEmbedding } from "./embedding-emulator.js";
import {
  consolidateProfile,
  findUsersDueForConsolidation,
  PROFILE_CONSOLIDATION_INTERVAL_MS,
  PROFILE_CONSOLIDATION_LOOKBACK_MS,
} from "./profile-consolidation.js";

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

const signupEnv = { signupsEnabled: "true", signupAllowlist: "" };
const HOUR = 60 * 60_000;

function ago(ms: number) {
  return new Date(Date.now() - ms);
}

function runtimeReplying(reply: string) {
  const requests: AgentRunRequest[] = [];
  const runtime: AgentRuntime = {
    describe: () => ({
      id: "test",
      contractVersion: "1",
      adapterVersion: "0.0.0",
      capabilities: { streaming: true, compaction: true, tools: true, scripted: false },
    }),
    async *run(request) {
      requests.push(request);
      yield { type: "done" as const, text: reply };
    },
    abort: async () => undefined,
  };
  return { runtime, requests };
}

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

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

  async function addMessages(
    threadId: string,
    messages: Array<{ role: "user" | "bot"; text: string; at: Date }>,
  ) {
    const thread = await prisma.thread.findUniqueOrThrow({ where: { id: threadId } });
    await prisma.message.createMany({
      data: messages.map((message, index) => ({
        threadId,
        seq: thread.nextMessageSeq + index,
        role: message.role,
        blocks: [{ kind: "text", text: message.text }],
        createdAt: message.at,
      })),
    });
    await prisma.thread.update({
      where: { id: threadId },
      data: { nextMessageSeq: thread.nextMessageSeq + messages.length },
    });
  }

  async function addRun(
    user: { userId: string; spaceId: string; main: { id: string; threadId: string } },
    trigger: string,
    updatedAt: Date,
  ) {
    const task = await prisma.task.create({
      data: {
        spaceId: user.spaceId,
        botId: user.main.id,
        threadId: user.main.threadId,
        userId: user.userId,
        prompt: "hello",
        status: "completed",
      },
    });
    const run = await prisma.run.create({
      data: {
        spaceId: user.spaceId,
        botId: user.main.id,
        threadId: user.main.threadId,
        taskId: task.id,
        userId: user.userId,
        status: "completed",
        trigger,
      },
    });
    await prisma.$executeRaw`UPDATE "runs" SET "updatedAt" = ${updatedAt.toISOString()}::timestamp(3) WHERE "id" = ${run.id}`;
  }

  function consolidateDeps(
    runtime: AgentRuntime,
    embeddings?: EmulatorEmbeddingProvider,
  ): Parameters<typeof consolidateProfile>[0] & { enqueue: ReturnType<typeof vi.fn> } {
    const enqueue = vi.fn(async (_job: BackgroundJob) => undefined);
    return {
      enqueue,
      prisma,
      runtime,
      jobs: { enqueue, cancel: async () => undefined, close: async () => undefined },
      memoryProviders: {
        resolve: async () =>
          embeddings ? { provider: {} as never, defaultScope: "isolated", embeddings } : null,
      },
      resolveModel: async () => ({ provider: "openai-compatible", id: "qwen3", apiKey: undefined }),
    };
  }

  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("learns from the user's settled app conversations and then waits for new ones", async () => {
    const user = await newUser();
    await addMessages(user.main.threadId, [
      { role: "user", text: "I work night shifts as a nurse in Berlin.", at: ago(HOUR) },
      { role: "bot", text: "Noted. How do you sleep?", at: ago(HOUR - 60_000) },
      { role: "user", text: "Badly, I have asthma.", at: ago(HOUR - 120_000) },
      { role: "user", text: "One more thing right now.", at: ago(60_000) },
    ]);
    const conversation = await prisma.externalConversation.create({
      data: {
        provider: "telegram",
        workspaceId: "workspace",
        externalKey: `chat-${user.userId}`,
        conversationId: "chat",
        spaceId: user.spaceId,
        botId: user.main.id,
        userId: user.userId,
      },
    });
    const channel = await prisma.thread.create({
      data: {
        spaceId: user.spaceId,
        externalConversationId: conversation.id,
        userId: user.userId,
      },
    });
    await addMessages(channel.id, [
      { role: "user", text: "Hi, I am Bob and I fly planes.", at: ago(HOUR) },
    ]);
    await prisma.userFact.create({
      data: {
        userId: user.userId,
        category: "preferences",
        value: "Prefers short answers.",
        source: "onboarding",
      },
    });
    const { runtime, requests } = runtimeReplying(
      JSON.stringify({
        facts: [
          { category: "work", value: "Works night shifts as a nurse in Berlin.", confidence: 0.9 },
          { category: "health", value: "Has asthma.", confidence: 0.8 },
          { category: "preferences", value: "prefers short answers" },
        ],
        summary: "A night-shift nurse in Berlin who likes short answers.",
      }),
    );
    const deps = consolidateDeps(runtime);

    await consolidateProfile(deps, user.userId);

    expect(requests).toHaveLength(1);
    expect(requests[0]!.prompt).toContain("I work night shifts as a nurse in Berlin.");
    expect(requests[0]!.prompt).toContain("- [preferences] Prefers short answers.");
    expect(requests[0]!.prompt).not.toContain("Bob");
    expect(requests[0]!.prompt).not.toContain("One more thing right now.");

    const facts = await prisma.userFact.findMany({
      where: { userId: user.userId },
      orderBy: { category: "asc" },
      select: { category: true, value: true, source: true, status: true, sensitive: true },
    });
    expect(facts).toEqual([
      {
        category: "health",
        value: "Has asthma.",
        source: "consolidation",
        status: "pending",
        sensitive: true,
      },
      {
        category: "preferences",
        value: "Prefers short answers.",
        source: "onboarding",
        status: "active",
        sensitive: false,
      },
      {
        category: "work",
        value: "Works night shifts as a nurse in Berlin.",
        source: "consolidation",
        status: "active",
        sensitive: false,
      },
    ]);
    const profile = await prisma.userProfile.findUniqueOrThrow({ where: { userId: user.userId } });
    expect(profile.summary).toBe("A night-shift nurse in Berlin who likes short answers.");
    expect(profile.consolidatedAt!.getTime()).toBeLessThan(ago(4 * 60_000).getTime());
    expect(profile.consolidatedAt!.getTime()).toBeGreaterThan(ago(HOUR).getTime());

    // Nothing new has settled since, so a second run reads nothing and asks no model.
    await consolidateProfile(deps, user.userId);
    expect(requests).toHaveLength(1);
    expect(await prisma.userFact.count({ where: { userId: user.userId } })).toBe(3);
  });

  it("stores fact vectors and skips restatements when embeddings are configured", async () => {
    const user = await newUser();
    await addMessages(user.main.threadId, [
      { role: "user", text: "Still nursing in Berlin. Started bouldering too.", at: ago(HOUR) },
    ]);
    await prisma.userFact.create({
      data: {
        userId: user.userId,
        category: "work",
        value: "Works as a nurse in Berlin",
        source: "user",
      },
    });
    const { runtime } = runtimeReplying(
      JSON.stringify({
        facts: [
          { category: "work", value: "Works as a nurse in Berlin, Germany" },
          { category: "habits", value: "Goes bouldering." },
        ],
        summary: "",
      }),
    );

    await consolidateProfile(
      consolidateDeps(runtime, new EmulatorEmbeddingProvider("emulator-v1")),
      user.userId,
    );

    const facts = await prisma.userFact.findMany({
      where: { userId: user.userId },
      orderBy: { createdAt: "asc" },
      select: { value: true, embedding: true, embeddingModel: true },
    });
    expect(facts).toEqual([
      {
        value: "Works as a nurse in Berlin",
        embedding: emulatedEmbedding("Works as a nurse in Berlin"),
        embeddingModel: "emulator-v1",
      },
      {
        value: "Goes bouldering.",
        embedding: emulatedEmbedding("Goes bouldering."),
        embeddingModel: "emulator-v1",
      },
    ]);
  });

  it("never pushes the profile past its fact limit", async () => {
    const user = await newUser();
    await addMessages(user.main.threadId, [
      { role: "user", text: "I run, swim and cycle every week.", at: ago(HOUR) },
    ]);
    await prisma.userFact.createMany({
      data: Array.from({ length: USER_FACT_LIMIT - 1 }, (_, index) => ({
        userId: user.userId,
        category: "habits",
        value: `Known habit ${index}`,
        source: "user",
      })),
    });
    const { runtime } = runtimeReplying(
      JSON.stringify({
        facts: [
          { category: "habits", value: "Runs every week." },
          { category: "habits", value: "Swims every week." },
          { category: "habits", value: "Cycles every week." },
        ],
        summary: "",
      }),
    );

    await consolidateProfile(consolidateDeps(runtime), user.userId);

    expect(await prisma.userFact.count({ where: { userId: user.userId } })).toBe(USER_FACT_LIMIT);
    expect(
      await prisma.userFact.count({ where: { userId: user.userId, source: "consolidation" } }),
    ).toBe(1);
    const profile = await prisma.userProfile.findUniqueOrThrow({ where: { userId: user.userId } });
    expect(profile.consolidatedAt).not.toBeNull();
  });

  it("schedules users with new user activity whose last consolidation is a day old", async () => {
    const now = new Date();
    const fresh = await newUser();
    await addRun(fresh, "user", ago(HOUR));

    const recent = await newUser();
    await addRun(recent, "user", ago(HOUR));
    await prisma.userProfile.create({
      data: { userId: recent.userId, consolidatedAt: ago(2 * HOUR) },
    });

    const returning = await newUser();
    await addRun(returning, "user", ago(HOUR));
    await prisma.userProfile.create({
      data: {
        userId: returning.userId,
        consolidatedAt: ago(PROFILE_CONSOLIDATION_INTERVAL_MS + HOUR),
      },
    });

    const idle = await newUser();
    await addRun(idle, "user", ago(3 * PROFILE_CONSOLIDATION_INTERVAL_MS));
    await prisma.userProfile.create({
      data: { userId: idle.userId, consolidatedAt: ago(2 * PROFILE_CONSOLIDATION_INTERVAL_MS) },
    });

    const routinesOnly = await newUser();
    await addRun(routinesOnly, "routine", ago(HOUR));

    const longGone = await newUser();
    await addRun(longGone, "user", ago(PROFILE_CONSOLIDATION_LOOKBACK_MS + HOUR));

    const due = await findUsersDueForConsolidation(prisma, { now, afterUserId: "", limit: 500 });
    expect(due).toEqual(expect.arrayContaining([fresh.userId, returning.userId]));
    for (const user of [recent, idle, routinesOnly, longGone]) {
      expect(due).not.toContain(user.userId);
    }

    const [first] = [fresh.userId, returning.userId].sort();
    const page = await findUsersDueForConsolidation(prisma, {
      now,
      afterUserId: first!,
      limit: 500,
    });
    expect(page).not.toContain(first);
    expect(page.every((userId) => userId > first!)).toBe(true);
  });
});