import { logger } from "@zerobyte/core/node"; import { serverEvents } from "~/server/core/events"; import { toMessage } from "~/server/utils/errors"; import type { ParsedTask, TaskResult } from "~/schemas/tasks"; import { taskStore } from "./tasks.store"; export class TaskCancelledError extends Error { readonly name = "TaskCancelledError"; readonly result: TResult | null; constructor(message: string, result: TResult | null = null) { super(message); this.result = result; } } type TaskLifecycleOptions = { taskId: string; label: string; cancellable?: boolean; prepare?: (signal: AbortSignal) => Promise<() => void>; run: (signal: AbortSignal) => Promise; onStarted?: (task: ParsedTask) => void | Promise; onSucceeded?: (task: ParsedTask, result: TResult) => void; onFailed?: (task: ParsedTask, errorMessage: string) => void; onCancelled?: (task: ParsedTask, errorMessage: string, result: TResult | null) => void; }; const abortControllers = new Map(); const emitTaskLifecycleEvent = (eventName: "task:started" | "task:finished", task: ParsedTask) => { serverEvents.emit(eventName, { organizationId: task.organizationId, taskId: task.id, kind: task.kind, resourceType: task.resourceType, resourceId: task.resourceId, status: task.status, }); }; const failTask = (options: TaskLifecycleOptions, errorMessage: string) => { try { const failedTask = taskStore.fail(options.taskId, errorMessage); options.onFailed?.(failedTask, errorMessage); return failedTask; } catch (error) { logger.warn(`Failed to fail ${options.label} ${options.taskId}: ${toMessage(error)}`); return null; } }; const cancelTask = ( options: TaskLifecycleOptions, errorMessage: string, result: TResult | null, ) => { try { const cancelledTask = taskStore.cancel(options.taskId, errorMessage, result); options.onCancelled?.(cancelledTask, errorMessage, result); return cancelledTask; } catch (error) { logger.warn(`Failed to cancel ${options.label} ${options.taskId}: ${toMessage(error)}`); return null; } }; const isTaskCancelledError = (error: unknown): error is TaskCancelledError => { return error instanceof TaskCancelledError; }; const isAbortError = (error: unknown) => { if (isTaskCancelledError(error)) { return true; } if (error instanceof DOMException && error.name === "AbortError") { return true; } return error instanceof Error && error.name === "AbortError"; }; export const requestTaskCancel = (taskId: string) => { const abortController = abortControllers.get(taskId); if (!abortController) { return false; } try { taskStore.requestCancel(taskId); } catch { return false; } abortController.abort(); return true; }; export const runTaskLifecycle = async (options: TaskLifecycleOptions) => { const abortController = new AbortController(); let cleanup: (() => void) | undefined; if (options.cancellable) { abortControllers.set(options.taskId, abortController); } try { cleanup = await options.prepare?.(abortController.signal); abortController.signal.throwIfAborted(); const startedTask = taskStore.markRunning(options.taskId); await options.onStarted?.(startedTask); emitTaskLifecycleEvent("task:started", startedTask); if (startedTask.cancellationRequested) { abortController.abort(); } const result = await options.run(abortController.signal); const completedTask = taskStore.complete(options.taskId, result); try { options.onSucceeded?.(completedTask, result); } catch (error) { logger.warn(`Failed to handle successful ${options.label} ${options.taskId}: ${toMessage(error)}`); } emitTaskLifecycleEvent("task:finished", completedTask); } catch (error) { if (abortController.signal.aborted || isAbortError(error)) { const cancelledError = isTaskCancelledError(error) ? error : null; const errorMessage = cancelledError?.message || toMessage(error) || "Task was cancelled"; const result = (cancelledError?.result as TResult | null) ?? null; const cancelledTask = cancelTask(options, errorMessage, result); if (cancelledTask) { emitTaskLifecycleEvent("task:finished", cancelledTask); } return; } const failedTask = failTask(options, toMessage(error)); if (failedTask) { emitTaskLifecycleEvent("task:finished", failedTask); } } finally { cleanup?.(); if (options.cancellable) { abortControllers.delete(options.taskId); } } };