import type { NotificationMessage, NotificationProvider } from "@rakazo/adapter-kit"; import { runContinueJob } from "@rakazo/adapter-kit"; import type { MessageBlock } from "@rakazo/contracts"; import { redactSecrets } from "@rakazo/core"; import type { Prisma, PrismaClient } from "@rakazo/db"; import { appendEventInTransaction, createThreadMessageInTransaction } from "@rakazo/db"; import { getLogger } from "@rakazo/logging"; import { withComputerScreenAvailability } from "./computer-screens.js"; import type { ExecutorDeps } from "./executor.js"; import { sanitizeTextBlocks } from "./prose-handoff-recovery.js"; import { TAKEOVER_RESUME_CHECKPOINTS, type TakeoverResumeCheckpoint } from "./takeover-resume.js"; // Helper used by publishMessage — also exported so executor.ts can call it directly export async function persistMessageInTransaction( tx: Prisma.TransactionClient, run: { id: string; spaceId: string; threadId: string; botId: string }, role: "user" | "bot" | "system", blocks: MessageBlock[], markUnread?: boolean, clientNonce?: string, ) { const message = await createThreadMessageInTransaction(tx, { threadId: run.threadId, role, blocks, botId: run.botId, runId: run.id, markUnread, clientNonce, }); const event = await appendEventInTransaction(tx, { spaceId: run.spaceId, threadId: run.threadId, botId: run.botId, type: "thread.message.created", runId: run.id, payload: { messageId: message.id, role, blocks }, }); return { message, eventSeq: event.seq }; } export async function computerScreenToolResult( work: () => Promise, finish?: (result: unknown) => Promise, ) { const result = await withComputerScreenAvailability(work); return finish ? finish(result) : result; } export async function runNotificationsEnabled( prisma: PrismaClient, run: { spaceId: string; userId: string; botId: string; threadId: string }, ): Promise { const source = await prisma.run.findFirst({ where: { botId: run.botId, threadId: run.threadId, spaceId: run.spaceId, userId: run.userId, }, select: { bot: { select: { notifyOnFinish: true } }, thread: { select: { groupId: true } }, }, }); return Boolean(source && (source.thread.groupId || source.bot.notifyOnFinish)); } export async function notifyRun( deps: ExecutorDeps, run: { spaceId: string; userId: string; botId: string; threadId: string }, message: NotificationMessage, ) { if (!deps.notifications) return; const enabled = await runNotificationsEnabled(deps.prisma, run).catch((error) => { getLogger().error("notification preference lookup", error); return false; }); if (!enabled) return; await deps.notifications .send(message, { operationId: "notify", traceId: run.botId, spaceId: run.spaceId, userId: run.userId, botId: run.botId, signal: new AbortController().signal, }) .catch((error) => { getLogger().error("run notification", error); }); } export async function renewRunLease( deps: ExecutorDeps, runId: string, workerId: string, fence: number, ): Promise { const renewed = await deps.prisma.run.updateMany({ where: { id: runId, status: "running", leaseOwner: workerId, leaseFence: fence }, data: { leaseExpiresAt: new Date(Date.now() + 5 * 60_000) }, }); return renewed.count === 1; } export function computerRetryDelay(fence: number): number { return Math.min(10_000, 250 * 2 ** Math.min(Math.max(fence - 1, 0), 5)); } export function computerRunRequeueData( resumeCheckpoint: TakeoverResumeCheckpoint | null, error: string | null = null, heldForTakeover = false, ) { return { status: heldForTakeover && !resumeCheckpoint ? ("waiting_takeover" as const) : ("queued" as const), error, leaseOwner: null, leaseExpiresAt: null, checkpoint: resumeCheckpoint, }; } export async function writeComputerRunRequeue( deps: ExecutorDeps, runId: string, workerId: string, fence: number, resumeCheckpoint: TakeoverResumeCheckpoint | null, heldForTakeover = false, error: string | null = null, ): Promise { const whereLease = { id: runId, status: "running" as const, leaseOwner: workerId, leaseFence: fence, }; const releasedHold = { status: "queued" as const, error, leaseOwner: null, leaseExpiresAt: null, }; const preserve = await deps.prisma.run.updateMany({ where: { ...whereLease, checkpoint: { in: [...TAKEOVER_RESUME_CHECKPOINTS] }, }, data: releasedHold, }); if (preserve.count === 1) return true; const planned = await deps.prisma.run.updateMany({ where: { ...whereLease, checkpoint: null }, data: computerRunRequeueData(resumeCheckpoint, error, heldForTakeover), }); if (planned.count === 1) return true; const retried = await deps.prisma.run.updateMany({ where: { ...whereLease, checkpoint: { in: [...TAKEOVER_RESUME_CHECKPOINTS] }, }, data: releasedHold, }); return retried.count === 1; } export async function requeueComputerRun( deps: ExecutorDeps, runId: string, workerId: string, fence: number, resumeCheckpoint: TakeoverResumeCheckpoint | null, heldForTakeover = false, ): Promise { const released = await writeComputerRunRequeue( deps, runId, workerId, fence, resumeCheckpoint, heldForTakeover, ); if (!released) return; await deps.jobs.enqueue({ ...runContinueJob(runId), availableAt: new Date(Date.now() + computerRetryDelay(fence)), }); } export function redactBlocks(blocks: MessageBlock[], secrets: string[]): MessageBlock[] { return sanitizeTextBlocks( blocks.map((block) => { if (block.kind === "text") { return { kind: "text" as const, text: redactSecrets(block.text, secrets) }; } if (block.kind === "bot_message_sent" || block.kind === "bot_message_received") { return { ...block, text: redactSecrets(block.text, secrets) }; } return block; }), ); } export async function publishMessage( deps: ExecutorDeps, run: { id: string; spaceId: string; threadId: string; botId: string }, role: "user" | "bot" | "system", blocks: MessageBlock[], markUnread?: boolean, clientNonce?: string, ) { const committed = await deps.prisma.$transaction((tx) => persistMessageInTransaction(tx, run, role, blocks, markUnread, clientNonce), ); await deps.events.notify(run.threadId, committed.eventSeq).catch((error) => { getLogger().error("thread message realtime notification", error); }); return committed.message; }