evOo-local/packages/adapters/src/delegated-tasks.ts

import { runContinueJob } from "@rakazo/adapter-kit";
import type { DelegatedTaskStatus, MessageBlock } from "@rakazo/contracts";
import {
  BOT_MESSAGE_MAX_LENGTH,
  botMessageHopExhausted,
  buildDelegatedTaskWakePrompt,
  clampBotMessage,
  DELEGATED_TASK_TITLE_MAX_LENGTH,
  DELEGATED_TASK_TRIGGER,
  delegatedTaskStatusForRun,
  isOpenDelegatedTaskStatus,
  nextBotMessageHop,
  renderDelegatedTaskInstruction,
  resolveBotAddress,
  taskBlockOf,
  withTaskStatus,
} from "@rakazo/core";
import type { Prisma, PrismaClient } from "@rakazo/db";
import {
  appendEventInTransaction,
  createThreadMessageInTransaction,
  ensureMainConversation,
  joinMainConversation,
  touchGroupUpdatedAt,
  withTransactionRetry,
} from "@rakazo/db";
import { getLogger } from "@rakazo/logging";
import type { ExecutorDeps } from "./executor-types.js";
import { botRunOutcomeText } from "./job-reconciler.js";

type DelegationDeps = Pick<ExecutorDeps, "prisma" | "events" | "jobs">;
type ThreadSeq = { threadId: string; seq: number };

const RECONCILE_BATCH_SIZE = 100;

function asBlocks(value: unknown): MessageBlock[] {
  return Array.isArray(value) ? (value as MessageBlock[]) : [];
}

async function lockThreads(tx: Prisma.TransactionClient, threadIds: string[]) {
  // One order everywhere so two writers cannot deadlock on the same pair of threads.
  for (const threadId of [...new Set(threadIds)].sort()) {
    await tx.$queryRaw`SELECT id FROM threads WHERE id = ${threadId} FOR UPDATE`;
  }
}

async function notifyThreads(deps: DelegationDeps, seqs: ThreadSeq[]) {
  for (const { threadId, seq } of seqs) {
    await deps.events.notify(threadId, seq).catch((error) => {
      getLogger().error("delegated task realtime notification", error);
    });
  }
}

/**
 * Hands work from the main bot to a specialist as a tracked task. The card is posted in the main
 * conversation; a specialist of the same space joins it and works there, one of another space
 * works in its own chat. Replays of the same tool call return the task created first.
 */
export async function delegateTask(
  deps: DelegationDeps,
  run: {
    id: string;
    botId: string;
    userId: string;
    sourceMessageId?: string | null;
  },
  sender: { id: string; name: string },
  input: { bot: string; title: string; brief: string; deadline?: string; deliveryKey: string },
): Promise<
  { ok: true; taskId: string; botId: string; name: string; note: string } | { error: string }
> {
  const title = input.title.trim();
  const brief = input.brief.trim();
  const requested = input.bot.trim();
  if (!requested) return { error: "bot is required" };
  if (!title) return { error: "title is required" };
  if (title.length > DELEGATED_TASK_TITLE_MAX_LENGTH) {
    return { error: `title exceeds the ${DELEGATED_TASK_TITLE_MAX_LENGTH} character limit` };
  }
  if (!brief) return { error: "brief is required" };
  if (brief.length > BOT_MESSAGE_MAX_LENGTH) {
    return { error: `brief exceeds the ${BOT_MESSAGE_MAX_LENGTH} character limit` };
  }
  let deadlineAt: Date | null = null;
  if (input.deadline?.trim()) {
    deadlineAt = new Date(input.deadline.trim());
    if (Number.isNaN(deadlineAt.getTime())) {
      return { error: "deadline must be an ISO 8601 date-time" };
    }
    if (deadlineAt.getTime() <= Date.now()) return { error: "deadline must be in the future" };
  }

  // Resolving the conversation also settles who the main bot is; only its lead delegates.
  const conversation = await ensureMainConversation(deps.prisma, run.userId);
  if (conversation.leadBotId !== sender.id) {
    return { error: "only the main agent can delegate tasks" };
  }

  const source = run.sourceMessageId
    ? await deps.prisma.message.findUnique({
        where: { id: run.sourceMessageId },
        select: { blocks: true },
      })
    : null;
  // A run woken by a finished task carries that task's card; its hop bounds chained delegation.
  const hop = nextBotMessageHop(taskBlockOf(asBlocks(source?.blocks))?.hop);
  if (botMessageHopExhausted(hop)) {
    return {
      error:
        "delegation limit reached for this chain; report back to the user instead of delegating again",
    };
  }

  const candidates = await deps.prisma.bot.findMany({
    where: {
      userId: run.userId,
      archivedAt: null,
      id: { not: sender.id },
      thread: { isNot: null },
      space: { memberships: { some: { userId: run.userId } } },
    },
    select: { id: true, name: true, spaceId: true, thread: { select: { id: true } } },
    orderBy: { createdAt: "asc" },
  });
  // A name shared across spaces resolves to the specialist next to the main conversation.
  const ordered = [
    ...candidates.filter((bot) => bot.spaceId === conversation.spaceId),
    ...candidates.filter((bot) => bot.spaceId !== conversation.spaceId),
  ];
  const target =
    ordered.find((bot) => bot.id === requested) ?? resolveBotAddress(ordered, { name: requested });
  if (!target?.thread) {
    if (requested === sender.id || requested.toLowerCase() === sender.name.toLowerCase()) {
      return { error: "you cannot delegate to yourself" };
    }
    return { error: "no specialist found with that name or id" };
  }
  const targetThreadId = target.thread.id;
  const inConversation = target.spaceId === conversation.spaceId;
  if (inConversation) {
    await joinMainConversation(deps.prisma, { userId: run.userId, botIds: [target.id] });
  }

  const deliveryKey = `delegated-task:${input.deliveryKey}`;
  const committed = await withTransactionRetry(() =>
    deps.prisma.$transaction(async (tx) => {
      await lockThreads(
        tx,
        inConversation ? [conversation.threadId] : [conversation.threadId, targetThreadId],
      );
      const replayed = await tx.message.findUnique({
        where: {
          threadId_clientNonce: { threadId: conversation.threadId, clientNonce: deliveryKey },
        },
        select: { id: true },
      });
      if (replayed) {
        const task = await tx.delegatedTask.findFirst({
          where: { messageId: replayed.id },
          select: { id: true },
        });
        if (!task) return { outcome: "error", error: "recorded delegation is incomplete" } as const;
        return { outcome: "replayed", taskId: task.id } as const;
      }
      const [senderRunning, assignee] = await Promise.all([
        tx.run.findFirst({
          where: { id: run.id, botId: run.botId, userId: run.userId, status: "running" },
          select: { id: true },
        }),
        tx.bot.findFirst({
          where: { id: target.id, userId: run.userId, archivedAt: null },
          select: { id: true },
        }),
      ]);
      if (!senderRunning) {
        return { outcome: "error", error: "source run is no longer active" } as const;
      }
      if (!assignee) {
        return { outcome: "error", error: `${target.name} is no longer available` } as const;
      }

      const task = await tx.delegatedTask.create({
        data: {
          userId: run.userId,
          spaceId: conversation.spaceId,
          threadId: conversation.threadId,
          requesterBotId: sender.id,
          assigneeBotId: target.id,
          assigneeSpaceId: target.spaceId,
          title,
          brief,
          deadlineAt,
        },
      });
      const block: MessageBlock = {
        kind: "task",
        taskId: task.id,
        assigneeBotId: target.id,
        assigneeBotName: target.name,
        title,
        text: brief,
        status: "queued",
        hop,
      };
      const card = await createThreadMessageInTransaction(tx, {
        threadId: conversation.threadId,
        role: "bot",
        blocks: [block],
        botId: sender.id,
        clientNonce: deliveryKey,
      });
      const seqs: ThreadSeq[] = [];
      const cardEvent = await appendEventInTransaction(tx, {
        spaceId: conversation.spaceId,
        threadId: conversation.threadId,
        botId: sender.id,
        type: "thread.message.created",
        payload: { messageId: card.id, role: "bot", blocks: [block] },
      });
      seqs.push({ threadId: conversation.threadId, seq: cardEvent.seq });
      let sourceMessageId = card.id;
      if (!inConversation) {
        // The specialist's own chat shows what it was asked; it wakes on that message.
        const inbound = await createThreadMessageInTransaction(tx, {
          threadId: targetThreadId,
          role: "user",
          blocks: [block],
          markUnread: true,
        });
        const inboundEvent = await appendEventInTransaction(tx, {
          spaceId: target.spaceId,
          threadId: targetThreadId,
          botId: target.id,
          type: "thread.message.created",
          payload: { messageId: inbound.id, role: "user", blocks: [block] },
        });
        seqs.push({ threadId: targetThreadId, seq: inboundEvent.seq });
        sourceMessageId = inbound.id;
      }
      const workThreadId = inConversation ? conversation.threadId : targetThreadId;
      const workTask = await tx.task.create({
        data: {
          spaceId: target.spaceId,
          botId: target.id,
          threadId: workThreadId,
          userId: run.userId,
          prompt: brief,
          status: "queued",
        },
      });
      const workRun = await tx.run.create({
        data: {
          spaceId: target.spaceId,
          botId: target.id,
          threadId: workThreadId,
          taskId: workTask.id,
          userId: run.userId,
          status: "queued",
          trigger: DELEGATED_TASK_TRIGGER,
          sourceMessageId,
        },
        select: { id: true },
      });
      await tx.delegatedTask.update({
        where: { id: task.id },
        data: { runId: workRun.id, messageId: card.id },
      });
      await touchGroupUpdatedAt(tx, conversation.groupId);
      return { outcome: "created", taskId: task.id, runId: workRun.id, seqs } as const;
    }),
  );
  if (committed.outcome === "error") return { error: committed.error };
  const note = inConversation
    ? `${target.name} works on it in the main conversation. You are woken there with the result to close the task; continue independent work and do not poll.`
    : `${target.name} works on it in its own chat in another space. You are woken with the result in the main conversation to close the task; continue independent work and do not poll.`;
  if (committed.outcome === "replayed") {
    return { ok: true, taskId: committed.taskId, botId: target.id, name: target.name, note };
  }
  await notifyThreads(deps, committed.seqs);
  await deps.jobs.enqueue(runContinueJob(committed.runId)).catch((error) => {
    // The queued run is durable; the job reconciler repairs a missed wake.
    getLogger().error("delegated task enqueue", error);
  });
  return { ok: true, taskId: committed.taskId, botId: target.id, name: target.name, note };
}

/** The instruction a specialist's delegated run works from. */
export async function loadDelegatedTaskInstruction(
  prisma: PrismaClient,
  run: { id: string; trigger: string },
): Promise<string | undefined> {
  if (run.trigger !== DELEGATED_TASK_TRIGGER) return undefined;
  const task = await prisma.delegatedTask.findUnique({
    where: { runId: run.id },
    select: { title: true, brief: true, deadlineAt: true, requesterBotId: true },
  });
  if (!task) return undefined;
  const requester = await prisma.bot.findUnique({
    where: { id: task.requesterBotId },
    select: { name: true },
  });
  return renderDelegatedTaskInstruction({
    fromName: requester?.name ?? "",
    title: task.title,
    brief: task.brief,
    deadlineAt: task.deadlineAt,
  });
}

/** Mirrors a delegated run's state on its task after each transition; never throws into the run. */
export async function syncDelegatedTaskRun(
  deps: DelegationDeps,
  run: { id: string; trigger: string },
): Promise<void> {
  if (run.trigger !== DELEGATED_TASK_TRIGGER) return;
  await syncDelegatedTask(deps, run.id).catch((error) => {
    getLogger().error("delegated task sync", error);
  });
}

async function updateTaskCard(
  tx: Prisma.TransactionClient,
  card: {
    threadId: string;
    messageId: string | null;
    spaceId: string;
    taskId: string;
    status: DelegatedTaskStatus;
    /** Author of the update event when the card has no bot, e.g. an inbound user message. */
    fallbackBotId: string;
  },
): Promise<{ exists: boolean; seq?: number }> {
  if (!card.messageId) return { exists: false };
  const message = await tx.message.findFirst({
    where: { id: card.messageId, threadId: card.threadId },
    select: { id: true, role: true, botId: true, blocks: true },
  });
  if (!message) return { exists: false };
  const blocks = withTaskStatus(asBlocks(message.blocks), card.taskId, card.status);
  if (!blocks) return { exists: true };
  await tx.message.update({
    where: { id: message.id },
    data: { blocks: blocks as Prisma.InputJsonValue },
  });
  const event = await appendEventInTransaction(tx, {
    spaceId: card.spaceId,
    threadId: card.threadId,
    botId: message.botId ?? card.fallbackBotId,
    type: "thread.message.updated",
    payload: { messageId: message.id, role: message.role, blocks },
  });
  return { exists: true, seq: event.seq };
}

/**
 * Brings the task of delegated run `runId` in line with the run: running, waiting, or terminal.
 * The first writer of a terminal status updates the cards, reports a result from another space
 * into the main conversation, and wakes the main bot there to close the task.
 */
export async function syncDelegatedTask(deps: DelegationDeps, runId: string): Promise<void> {
  const committed = await withTransactionRetry(() =>
    deps.prisma.$transaction(async (tx) => {
      const task = await tx.delegatedTask.findUnique({ where: { runId } });
      if (!task || !isOpenDelegatedTaskStatus(task.status)) return null;
      const run = await tx.run.findUnique({
        where: { id: runId },
        select: { status: true, error: true, threadId: true, spaceId: true, sourceMessageId: true },
      });
      // A run deleted with its bot can never report back.
      const status = run ? delegatedTaskStatusForRun(run.status) : "cancelled";
      if (!status || status === task.status) return null;
      await lockThreads(tx, run ? [task.threadId, run.threadId] : [task.threadId]);
      const terminal = !isOpenDelegatedTaskStatus(status);
      const written = status === "done" ? (await botRunOutcomeText(tx, runId)).text : "";
      const result =
        status === "done"
          ? written || "The specialist finished without a written result."
          : status === "failed"
            ? clampBotMessage(run?.error ?? "unknown error")
            : null;
      const updated = await tx.delegatedTask.updateMany({
        where: { id: task.id, status: task.status },
        data: { status, ...(terminal ? { result, completedAt: new Date() } : {}) },
      });
      if (updated.count !== 1) return null;

      const seqs: ThreadSeq[] = [];
      const card = await updateTaskCard(tx, {
        threadId: task.threadId,
        messageId: task.messageId,
        spaceId: task.spaceId,
        taskId: task.id,
        status,
        fallbackBotId: task.requesterBotId,
      });
      if (card.seq !== undefined) seqs.push({ threadId: task.threadId, seq: card.seq });
      const crossSpace = Boolean(run && run.threadId !== task.threadId);
      if (run && crossSpace) {
        const inbound = await updateTaskCard(tx, {
          threadId: run.threadId,
          messageId: run.sourceMessageId,
          spaceId: run.spaceId,
          taskId: task.id,
          status,
          fallbackBotId: task.assigneeBotId,
        });
        if (inbound.seq !== undefined) seqs.push({ threadId: run.threadId, seq: inbound.seq });
      }
      if ((status !== "done" && status !== "failed") || result === null) return { seqs };

      const [assignee, requester] = await Promise.all([
        tx.bot.findUnique({ where: { id: task.assigneeBotId }, select: { name: true } }),
        tx.bot.findFirst({
          where: { id: task.requesterBotId, archivedAt: null },
          select: { id: true },
        }),
      ]);
      if (written && crossSpace) {
        // The work happened in the specialist's own space; its result lands where the task lives.
        const blocks: MessageBlock[] = [{ kind: "text", text: clampBotMessage(written) }];
        const reported = await createThreadMessageInTransaction(tx, {
          threadId: task.threadId,
          role: "bot",
          blocks,
          botId: task.assigneeBotId,
        });
        const event = await appendEventInTransaction(tx, {
          spaceId: task.spaceId,
          threadId: task.threadId,
          botId: task.assigneeBotId,
          type: "thread.message.created",
          payload: { messageId: reported.id, role: "bot", blocks },
        });
        seqs.push({ threadId: task.threadId, seq: event.seq });
      }
      if (!requester) return { seqs };
      const prompt = buildDelegatedTaskWakePrompt({
        assigneeName: assignee?.name ?? "",
        title: task.title,
        status,
        result,
      });
      const wakeTask = await tx.task.create({
        data: {
          spaceId: task.spaceId,
          botId: requester.id,
          threadId: task.threadId,
          userId: task.userId,
          prompt,
          status: "queued",
        },
      });
      const wakeRun = await tx.run.create({
        data: {
          spaceId: task.spaceId,
          botId: requester.id,
          threadId: task.threadId,
          taskId: wakeTask.id,
          userId: task.userId,
          status: "queued",
          // Closing the task is part of the delegation; it has no task of its own.
          trigger: DELEGATED_TASK_TRIGGER,
          // The card links the wake to the task, which bounds chained delegation.
          sourceMessageId: card.exists ? task.messageId : null,
        },
        select: { id: true },
      });
      return { seqs, wakeRunId: wakeRun.id };
    }),
  );
  if (!committed) return;
  await notifyThreads(deps, committed.seqs);
  if (committed.wakeRunId) {
    await deps.jobs.enqueue(runContinueJob(committed.wakeRunId)).catch((error) => {
      getLogger().error("delegated task wake enqueue", error);
    });
  }
}

/** Settles open tasks whose run ended without the executor reporting it (stop, clear, crash). */
export async function reconcileDelegatedTasks(deps: DelegationDeps): Promise<void> {
  const stale = await deps.prisma.$queryRaw<Array<{ runId: string }>>`
    SELECT t."runId"
    FROM delegated_tasks t
    LEFT JOIN runs r ON r.id = t."runId"
    WHERE t."runId" IS NOT NULL
      AND t.status IN ('queued', 'running', 'waiting')
      AND (r.id IS NULL OR r.status IN ('completed', 'failed', 'cancelled'))
    ORDER BY t."updatedAt" ASC
    LIMIT ${RECONCILE_BATCH_SIZE}
  `;
  for (const { runId } of stale) {
    await syncDelegatedTask(deps, runId).catch((error) => {
      getLogger().error("delegated task reconciliation", error);
    });
  }
}