evOo-local/packages/testkit/src/delegation.postgres.test.ts

import { mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import type { JobPublisher } from "@rakazo/adapter-kit";
import { ComposioEmulator, delegateTask, reconcileDelegatedTasks } from "@rakazo/adapters";
import type { MessageBlock } from "@rakazo/contracts";
import { createThreadEvents, ensureMainConversation, MAIN_BOT_SPAWN_KEY } from "@rakazo/db";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import type { createApp } from "../../../apps/api/src/app.ts";
import { sessionCookieHeader } from "./index.js";

type App = { request: (input: string, init?: RequestInit) => Response | Promise<Response> };
type AppHandles = Awaited<ReturnType<typeof createApp>>;

process.env.WAKEUP_DRIVER = "memory";
process.env.SANDBOX_PROVIDER = "fake";
process.env.AGENT_RUNTIME = "scripted";

const hasDb = process.env.VERIFY_DATABASE === "1" && Boolean(process.env.DATABASE_URL);
const describeWithDatabase = hasDb ? describe.sequential : describe.skip;

const idleJobs: JobPublisher = {
  enqueue: async () => undefined,
  cancel: async () => undefined,
  close: async () => undefined,
};

describeWithDatabase("delegated tasks (PostgreSQL)", () => {
  let handles: AppHandles;
  let app: App;
  const stamp = `${Date.now()}-${Math.random().toString(36).slice(2)}`;
  const dataDir = mkdtempSync(path.join(tmpdir(), "rakazo-delegation-"));

  beforeAll(async () => {
    const { createApp } = await import("../../../apps/api/src/app.ts");
    handles = await createApp({
      databaseUrl: process.env.DATABASE_URL!,
      dataDir,
      sandboxProvider: "fake",
      agentRuntime: "scripted",
      wakeupDriver: "memory",
      signupsEnabled: "true",
      composio: new ComposioEmulator(),
    });
    app = handles.app;
  });

  afterAll(async () => {
    await handles?.stop();
    rmSync(dataDir, { recursive: true, force: true });
  });

  /** A user whose onboarding bot became the main bot, plus their main conversation. */
  async function userWithMainBot(label: string) {
    const cookie = await signup(app, `${label}-${stamp}@rakazo.test`, label);
    const me = await rpc<{ userId: string; spaceId: string }>(app, cookie, "me");
    await rpc(app, cookie, "bots/create", {
      ...botInput("Chief"),
      spawnKey: MAIN_BOT_SPAWN_KEY,
    });
    const conversation = await ensureMainConversation(handles.prisma, me.userId);
    return { cookie, me, conversation };
  }

  async function waitForTask(userId: string, title: string, statuses: string[]) {
    let task: Awaited<ReturnType<typeof handles.prisma.delegatedTask.findFirst>> = null;
    await waitForDatabase(async () => {
      task = await handles.prisma.delegatedTask.findFirst({ where: { userId, title } });
      return Boolean(task && statuses.includes(task.status));
    });
    return task!;
  }

  async function waitForWakeRun(taskThreadId: string, botId: string, after: Date) {
    let wake: { id: string; status: string; trigger: string } | null = null;
    await waitForDatabase(async () => {
      wake = await handles.prisma.run.findFirst({
        where: { threadId: taskThreadId, botId, trigger: "delegation", createdAt: { gte: after } },
        select: { id: true, status: true, trigger: true },
        orderBy: { createdAt: "desc" },
      });
      return wake?.status === "completed";
    });
    return wake!;
  }

  /** A bot turn in its own chat that is mid-run, as a tool call sees it. */
  async function runningRun(me: { userId: string; spaceId: string }, botId: string) {
    const thread = await handles.prisma.thread.findUniqueOrThrow({ where: { botId } });
    const task = await handles.prisma.task.create({
      data: {
        spaceId: me.spaceId,
        botId,
        threadId: thread.id,
        userId: me.userId,
        prompt: "delegate",
        status: "running",
      },
    });
    return handles.prisma.run.create({
      data: {
        spaceId: me.spaceId,
        botId,
        threadId: thread.id,
        taskId: task.id,
        userId: me.userId,
        status: "running",
        trigger: "user",
      },
    });
  }

  function taskBlock(blocks: unknown) {
    return (blocks as MessageBlock[]).find(
      (block): block is Extract<MessageBlock, { kind: "task" }> => block.kind === "task",
    );
  }

  it("runs a same-space task in the main conversation and wakes the main bot to close it", async () => {
    const { cookie, me, conversation } = await userWithMainBot("delegate-same");
    const researcher = await rpc<{ id: string }>(
      app,
      cookie,
      "bots/create",
      botInput("Researcher"),
    );
    const startedAt = new Date();

    const sent = await rpc<{ runId: string }>(app, cookie, "threads/send", {
      botId: conversation.leadBotId,
      text: "delegate to the bot named Researcher: collect three owl facts",
    });
    const task = await waitForTask(me.userId, "collect three owl facts", ["done"]);
    expect(task).toMatchObject({
      requesterBotId: conversation.leadBotId,
      assigneeBotId: researcher.id,
      spaceId: conversation.spaceId,
      assigneeSpaceId: me.spaceId,
      threadId: conversation.threadId,
      brief: "collect three owl facts",
    });
    expect(task.result).toContain("done. i handled: collect three owl facts");
    expect(task.completedAt).toBeInstanceOf(Date);

    // The specialist joined the stage and worked in the main conversation, visibly.
    expect(
      await handles.prisma.chatGroupMember.findFirst({
        where: { groupId: conversation.groupId, botId: researcher.id },
      }),
    ).not.toBeNull();
    const work = await handles.prisma.run.findUniqueOrThrow({ where: { id: task.runId! } });
    expect(work).toMatchObject({
      botId: researcher.id,
      threadId: conversation.threadId,
      trigger: "delegation",
      status: "completed",
      sourceMessageId: task.messageId,
    });
    const reply = await handles.prisma.message.findFirst({
      where: { threadId: conversation.threadId, runId: work.id, botId: researcher.id },
    });
    expect(JSON.stringify(reply?.blocks)).toContain("done. i handled: collect three owl facts");

    // The card lives in the main conversation and followed the task to done.
    const card = await handles.prisma.message.findUniqueOrThrow({ where: { id: task.messageId! } });
    expect(card).toMatchObject({ threadId: conversation.threadId, botId: conversation.leadBotId });
    expect(taskBlock(card.blocks)).toMatchObject({
      taskId: task.id,
      assigneeBotName: "Researcher",
      title: "collect three owl facts",
      status: "done",
      hop: 1,
    });
    const updates = await handles.prisma.event.findMany({
      where: { threadId: conversation.threadId, type: "thread.message.updated" },
      orderBy: { seq: "asc" },
    });
    expect(
      updates
        .filter((event) => (event.payload as { messageId?: string }).messageId === card.id)
        .map((event) => taskBlock((event.payload as { blocks: unknown }).blocks)?.status),
    ).toEqual(["running", "done"]);

    // The main bot's own turn went on, and a wake closed the task in the main conversation.
    expect((await handles.prisma.run.findUniqueOrThrow({ where: { id: sent.runId } })).status).toBe(
      "completed",
    );
    const wake = await waitForWakeRun(conversation.threadId, conversation.leadBotId, startedAt);
    const closing = await handles.prisma.message.findFirst({ where: { runId: wake.id } });
    expect(JSON.stringify(closing?.blocks)).toContain("closing the task: collect three owl facts");

    // Snapshots, activity, and tasks.list all accept the card and the new trigger.
    const snapshot = await rpc<{ messages: Array<{ id: string; blocks: MessageBlock[] }> }>(
      app,
      cookie,
      "threads/get",
      { groupId: conversation.groupId },
    );
    expect(taskBlock(snapshot.messages.find((m) => m.id === card.id)?.blocks)?.status).toBe("done");
    const activity = await rpc<{ runs: Array<{ runId: string; trigger: string }> }>(
      app,
      cookie,
      "runs/list",
      { filter: "recent" },
    );
    expect(activity.runs).toContainEqual(
      expect.objectContaining({ runId: work.id, trigger: "delegation" }),
    );
    const listed = await rpc<Array<{ id: string; status: string; result: string | null }>>(
      app,
      cookie,
      "tasks/list",
      { threadId: conversation.threadId },
    );
    expect(listed).toEqual([
      expect.objectContaining({ id: task.id, status: "done", result: task.result }),
    ]);
    expect(await rpc(app, cookie, "tasks/list", { status: "running" })).toEqual([]);
  });

  it("delegates from a group run in the main conversation to another specialist", async () => {
    const { cookie, me, conversation } = await userWithMainBot("delegate-group");
    const researcher = await rpc<{ id: string }>(
      app,
      cookie,
      "bots/create",
      botInput("Researcher"),
    );
    const writer = await rpc<{ id: string }>(app, cookie, "bots/create", botInput("Writer"));
    // A second member lets the conversation take sends before the lead-only routing lands.
    await handles.prisma.chatGroupMember.create({
      data: { groupId: conversation.groupId, botId: researcher.id },
    });

    await rpc(app, cookie, "threads/send", {
      groupId: conversation.groupId,
      text: "delegate to the bot named Writer: draft a haiku about owls",
    });
    const task = await waitForTask(me.userId, "draft a haiku about owls", ["done"]);
    expect(task).toMatchObject({
      requesterBotId: conversation.leadBotId,
      assigneeBotId: writer.id,
      threadId: conversation.threadId,
    });
    const work = await handles.prisma.run.findUniqueOrThrow({ where: { id: task.runId! } });
    expect(work).toMatchObject({ threadId: conversation.threadId, trigger: "delegation" });
    const joined = await handles.prisma.event.findFirst({
      where: { threadId: conversation.threadId, type: "group.updated" },
      orderBy: { seq: "desc" },
    });
    expect(joined?.payload).toMatchObject({ addedBotIds: [writer.id] });
  });

  it("lets a project's own main agent delegate to the project's specialists in its conversation", async () => {
    const { cookie, me, conversation: home } = await userWithMainBot("delegate-project");
    const project = await rpc<{
      id: string;
      groups: Array<{ id: string; kind: string; members: Array<{ botId: string }> }>;
    }>(app, cookie, "spaces/create", { name: "Launch" });
    const main = project.groups.find((group) => group.kind === "main");
    expect(main).toBeDefined();
    expect(
      await rpc(app, cookie, "spaces/ensureMain", { spaceId: project.id }, project.id),
    ).toMatchObject({ groupId: main!.id, spaceId: project.id });
    const lead = main!.members[0]!.botId;
    expect(lead).not.toBe(home.leadBotId);
    const designer = await rpc<{ id: string }>(
      app,
      cookie,
      "bots/create",
      botInput("Designer"),
      project.id,
    );
    const thread = await handles.prisma.thread.findUniqueOrThrow({ where: { groupId: main!.id } });
    const startedAt = new Date();

    await rpc(
      app,
      cookie,
      "threads/send",
      { groupId: main!.id, text: "delegate to the bot named Designer: sketch a launch logo" },
      project.id,
    );
    const task = await waitForTask(me.userId, "sketch a launch logo", ["done"]);
    expect(task).toMatchObject({
      requesterBotId: lead,
      assigneeBotId: designer.id,
      spaceId: project.id,
      assigneeSpaceId: project.id,
      threadId: thread.id,
    });
    // The specialist joins the project's conversation, never the home one.
    expect(
      await handles.prisma.chatGroupMember.findFirst({
        where: { groupId: main!.id, botId: designer.id },
      }),
    ).not.toBeNull();
    expect(
      await handles.prisma.chatGroupMember.findFirst({
        where: { groupId: home.groupId, botId: designer.id },
      }),
    ).toBeNull();
    await waitForWakeRun(thread.id, lead, startedAt);
  });

  it("runs an other-space task in the specialist's own chat and reports the result back", async () => {
    const { cookie, me, conversation } = await userWithMainBot("delegate-cross");
    const lab = await rpc<{ id: string }>(app, cookie, "spaces/create", { name: "Lab" });
    const analyst = await rpc<{ id: string }>(
      app,
      cookie,
      "bots/create",
      botInput("Analyst"),
      lab.id,
    );
    const startedAt = new Date();

    await rpc(app, cookie, "threads/send", {
      botId: conversation.leadBotId,
      text: "delegate to the bot named Analyst: summarize the quarter",
    });
    const task = await waitForTask(me.userId, "summarize the quarter", ["done"]);
    expect(task).toMatchObject({
      assigneeBotId: analyst.id,
      assigneeSpaceId: lab.id,
      spaceId: conversation.spaceId,
      threadId: conversation.threadId,
    });
    const analystThread = await handles.prisma.thread.findUniqueOrThrow({
      where: { botId: analyst.id },
    });
    const work = await handles.prisma.run.findUniqueOrThrow({ where: { id: task.runId! } });
    expect(work).toMatchObject({
      spaceId: lab.id,
      threadId: analystThread.id,
      trigger: "delegation",
      status: "completed",
    });
    // The specialist's own chat shows what it was asked, with the same live status.
    const inbound = await handles.prisma.message.findUniqueOrThrow({
      where: { id: work.sourceMessageId! },
    });
    expect(inbound).toMatchObject({ threadId: analystThread.id, role: "user" });
    expect(taskBlock(inbound.blocks)).toMatchObject({ taskId: task.id, status: "done" });
    // It never joins a conversation outside its space; its result is reported there instead.
    expect(
      await handles.prisma.chatGroupMember.findFirst({
        where: { groupId: conversation.groupId, botId: analyst.id },
      }),
    ).toBeNull();
    const reported = await handles.prisma.message.findFirst({
      where: { threadId: conversation.threadId, botId: analyst.id },
    });
    expect(JSON.stringify(reported?.blocks)).toContain("done. i handled: summarize the quarter");
    await waitForWakeRun(conversation.threadId, conversation.leadBotId, startedAt);

    const fromLab = await rpc<Array<{ id: string }>>(app, cookie, "tasks/list", {}, lab.id);
    expect(fromLab.map((listed) => listed.id)).toContain(task.id);
  });

  it("reports a failed task to the main bot and leaves a stopped one cancelled", async () => {
    const { cookie, me, conversation } = await userWithMainBot("delegate-fail");
    await rpc(app, cookie, "bots/create", botInput("Researcher"));
    // The scripted runtime fails or hangs on phrases in its prompt, so these tasks start
    // directly from a main-bot turn instead of from a user message that holds the phrase.
    const mainRun = await runningRun(me, conversation.leadBotId);
    const main = { id: conversation.leadBotId, name: "Chief" };
    const deps = {
      prisma: handles.prisma,
      jobs: handles.jobs,
      events: createThreadEvents(handles.prisma),
    };
    const startedAt = new Date();

    await delegateTask(deps, mainRun, main, {
      bot: "Researcher",
      title: "Doomed",
      brief: "fail this run on purpose",
      deliveryKey: "fail",
    });
    const failed = await waitForTask(me.userId, "Doomed", ["failed"]);
    expect(failed.result).toContain("Scripted run failure");
    expect(
      taskBlock(
        (await handles.prisma.message.findUniqueOrThrow({ where: { id: failed.messageId! } }))
          .blocks,
      )?.status,
    ).toBe("failed");
    await waitForWakeRun(conversation.threadId, conversation.leadBotId, startedAt);

    await delegateTask(deps, mainRun, main, {
      bot: "Researcher",
      title: "Endless",
      brief: "keep working until i stop you",
      deliveryKey: "hang",
    });
    await waitForTask(me.userId, "Endless", ["running"]);
    const stoppedAt = new Date();
    await rpc(app, cookie, "threads/stop", { groupId: conversation.groupId });
    await reconcileDelegatedTasks(deps);
    const stopped = await waitForTask(me.userId, "Endless", ["cancelled"]);
    expect(stopped.result).toBeNull();
    expect(
      taskBlock(
        (await handles.prisma.message.findUniqueOrThrow({ where: { id: stopped.messageId! } }))
          .blocks,
      )?.status,
    ).toBe("cancelled");
    // Stopping is the user's own decision; it does not wake the main bot.
    expect(
      await handles.prisma.run.count({
        where: {
          threadId: conversation.threadId,
          botId: conversation.leadBotId,
          createdAt: { gte: stoppedAt },
        },
      }),
    ).toBe(0);
  });

  it("asks the user in the main conversation while the task waits, then resumes it", async () => {
    const { cookie, me, conversation } = await userWithMainBot("delegate-ask");
    const researcher = await rpc<{ id: string }>(
      app,
      cookie,
      "bots/create",
      botInput("Researcher"),
    );
    const mainRun = await runningRun(me, conversation.leadBotId);
    await delegateTask(
      { prisma: handles.prisma, jobs: handles.jobs, events: createThreadEvents(handles.prisma) },
      mainRun,
      { id: conversation.leadBotId, name: "Chief" },
      { bot: "Researcher", title: "City", brief: "ask me which city to use", deliveryKey: "ask" },
    );
    const waiting = await waitForTask(me.userId, "City", ["waiting"]);
    expect(
      taskBlock(
        (await handles.prisma.message.findUniqueOrThrow({ where: { id: waiting.messageId! } }))
          .blocks,
      )?.status,
    ).toBe("waiting");
    // The specialist's question is asked where the user follows the task, under its own name.
    const ask = await handles.prisma.message.findFirstOrThrow({
      where: { threadId: conversation.threadId, runId: waiting.runId!, botId: researcher.id },
      orderBy: { seq: "desc" },
    });
    expect(JSON.stringify(ask.blocks)).toContain("Which city should I use?");

    await rpc(app, cookie, "threads/answer", {
      groupId: conversation.groupId,
      runId: waiting.runId,
      messageId: ask.id,
      answer: "Paris",
    });
    const done = await waitForTask(me.userId, "City", ["done"]);
    expect(done.runId).toBe(waiting.runId);
    expect(done.result).toContain("Paris");
  });

  it("lets only the main bot delegate and replays a repeated tool call", async () => {
    const { cookie, me, conversation } = await userWithMainBot("delegate-rules");
    const researcher = await rpc<{ id: string; name: string }>(
      app,
      cookie,
      "bots/create",
      botInput("Researcher"),
    );
    const deps = {
      prisma: handles.prisma,
      jobs: idleJobs,
      events: createThreadEvents(handles.prisma),
    };
    const input = { bot: "Researcher", title: "Owl facts", brief: "Collect owl facts." };

    const specialistRun = await runningRun(me, researcher.id);
    expect(
      await delegateTask(deps, specialistRun, researcher, { ...input, deliveryKey: "k-1" }),
    ).toEqual({ error: "only the main agent can delegate tasks" });

    const mainRun = await runningRun(me, conversation.leadBotId);
    const main = { id: conversation.leadBotId, name: "Chief" };
    expect(
      await delegateTask(deps, mainRun, main, { ...input, bot: "Chief", deliveryKey: "k-2" }),
    ).toEqual({ error: "you cannot delegate to yourself" });
    expect(
      await delegateTask(deps, mainRun, main, { ...input, bot: "Nobody", deliveryKey: "k-3" }),
    ).toEqual({ error: "no specialist found with that name or id" });
    expect(
      await delegateTask(deps, mainRun, main, {
        ...input,
        deadline: "2001-01-01T00:00:00Z",
        deliveryKey: "k-4",
      }),
    ).toEqual({ error: "deadline must be in the future" });

    const first = await delegateTask(deps, mainRun, main, { ...input, deliveryKey: "k-5" });
    const again = await delegateTask(deps, mainRun, main, { ...input, deliveryKey: "k-5" });
    expect(first).toMatchObject({ ok: true, botId: researcher.id, name: "Researcher" });
    expect(again).toEqual(first);
    expect(await handles.prisma.delegatedTask.count({ where: { userId: me.userId } })).toBe(1);
  });
});

function botInput(name: string) {
  return { name, title: "", description: "", instructions: "", notifyOnFinish: false };
}

async function signup(app: App, email: string, name: string) {
  const res = await app.request("/api/auth/sign-up/email", {
    method: "POST",
    headers: { "content-type": "application/json", origin: "http://127.0.0.1:5173" },
    body: JSON.stringify({ email, password: "password12", name }),
  });
  if (res.status >= 400) throw new Error(`signup failed ${res.status}: ${await res.text()}`);
  return sessionCookieHeader(res);
}

async function rpc<T>(
  app: App,
  cookie: string,
  proc: string,
  body: unknown = {},
  spaceId?: string,
): Promise<T> {
  const res = await app.request(`/rpc/${proc}`, {
    method: "POST",
    headers: {
      "content-type": "application/json",
      cookie,
      origin: "http://127.0.0.1:5173",
      ...(spaceId ? { "x-rakazo-space-id": spaceId } : {}),
    },
    body: JSON.stringify({ json: body }),
  });
  const text = await res.text();
  const parsed = JSON.parse(text) as { json?: T; error?: { message?: string } };
  if (res.status >= 400 || parsed.error) {
    throw new Error(`${proc} ${res.status}: ${parsed.error?.message ?? text}`);
  }
  return parsed.json as T;
}

async function waitForDatabase(pred: () => Promise<boolean>, timeoutMs = 30_000) {
  const start = Date.now();
  while (Date.now() - start < timeoutMs) {
    if (await pred()) return;
    await new Promise((resolve) => setTimeout(resolve, 50));
  }
  throw new Error("timeout waiting for database state");
}