diff --git a/.changeset/weak-jobs-hide.md b/.changeset/weak-jobs-hide.md new file mode 100644 index 0000000000..0be1f49588 --- /dev/null +++ b/.changeset/weak-jobs-hide.md @@ -0,0 +1,7 @@ +--- +"@trigger.dev/sdk": patch +"trigger.dev": patch +"@trigger.dev/core": patch +--- + +v4: New lifecycle hooks diff --git a/.cursor/rules/executing-commands.mdc b/.cursor/rules/executing-commands.mdc index eac17379e4..0d36b44949 100644 --- a/.cursor/rules/executing-commands.mdc +++ b/.cursor/rules/executing-commands.mdc @@ -13,12 +13,12 @@ But often, when running tests, it's better to `cd` into the directory and then r ``` cd apps/webapp -pnpm run test +pnpm run test --run ``` This way you can run for a single file easily: ``` cd internal-packages/run-engine -pnpm run test ./src/engine/tests/ttl.test.ts +pnpm run test ./src/engine/tests/ttl.test.ts --run ``` diff --git a/apps/webapp/app/components/runs/v3/RunIcon.tsx b/apps/webapp/app/components/runs/v3/RunIcon.tsx index a557e5cd35..7c43fcfae6 100644 --- a/apps/webapp/app/components/runs/v3/RunIcon.tsx +++ b/apps/webapp/app/components/runs/v3/RunIcon.tsx @@ -44,6 +44,9 @@ export function RunIcon({ name, className, spanName }: TaskIconProps) { } if (!name) return ; + if (tablerIcons.has(name)) { + return ; + } switch (name) { case "task": @@ -73,6 +76,28 @@ export function RunIcon({ name, className, spanName }: TaskIconProps) { return ; case "fatal": return ; + case "task-middleware": + return ; + case "task-fn-run": + return ; + case "task-hook-init": + return ; + case "task-hook-onStart": + return ; + case "task-hook-onSuccess": + return ; + case "task-hook-onFailure": + return ; + case "task-hook-onComplete": + return ; + case "task-hook-onWait": + return ; + case "task-hook-onResume": + return ; + case "task-hook-catchError": + return ; + case "task-hook-cleanup": + return ; } return ; diff --git a/internal-packages/run-engine/src/engine/errors.ts b/internal-packages/run-engine/src/engine/errors.ts index 81e3b598bd..e1dd34eac4 100644 --- a/internal-packages/run-engine/src/engine/errors.ts +++ b/internal-packages/run-engine/src/engine/errors.ts @@ -13,6 +13,9 @@ export function runStatusFromError(error: TaskRunError): TaskRunStatus { //e.g. a bug switch (error.code) { case "RECURSIVE_WAIT_DEADLOCK": + case "TASK_INPUT_ERROR": + case "TASK_OUTPUT_ERROR": + case "TASK_MIDDLEWARE_ERROR": return "COMPLETED_WITH_ERRORS"; case "TASK_RUN_CANCELLED": return "CANCELED"; @@ -41,8 +44,6 @@ export function runStatusFromError(error: TaskRunError): TaskRunStatus { case "TASK_RUN_STALLED_EXECUTING_WITH_WAITPOINTS": case "TASK_HAS_N0_EXECUTION_SNAPSHOT": case "GRACEFUL_EXIT_TIMEOUT": - case "TASK_INPUT_ERROR": - case "TASK_OUTPUT_ERROR": case "POD_EVICTED": case "POD_UNKNOWN_ERROR": case "TASK_EXECUTION_ABORTED": diff --git a/packages/cli-v3/src/build/bundle.ts b/packages/cli-v3/src/build/bundle.ts index 0c2dfa5631..b326a5f769 100644 --- a/packages/cli-v3/src/build/bundle.ts +++ b/packages/cli-v3/src/build/bundle.ts @@ -47,6 +47,7 @@ export type BundleResult = { runControllerEntryPoint: string | undefined; indexWorkerEntryPoint: string | undefined; indexControllerEntryPoint: string | undefined; + initEntryPoint: string | undefined; stop: (() => Promise) | undefined; }; @@ -229,11 +230,26 @@ export async function getBundleResultFromBuild( let runControllerEntryPoint: string | undefined; let indexWorkerEntryPoint: string | undefined; let indexControllerEntryPoint: string | undefined; + let initEntryPoint: string | undefined; const configEntryPoint = resolvedConfig.configFile ? relative(resolvedConfig.workingDir, resolvedConfig.configFile) : "trigger.config.ts"; + // Check if the entry point is an init.ts file at the root of a trigger directory + function isInitEntryPoint(entryPoint: string): boolean { + const normalizedEntryPoint = entryPoint.replace(/\\/g, "/"); // Normalize path separators + const initFileNames = ["init.ts", "init.mts", "init.cts", "init.js", "init.mjs", "init.cjs"]; + + // Check if it's directly in one of the trigger directories + return resolvedConfig.dirs.some((dir) => { + const normalizedDir = dir.replace(/\\/g, "/"); + return initFileNames.some( + (fileName) => normalizedEntryPoint === `${normalizedDir}/${fileName}` + ); + }); + } + for (const [outputPath, outputMeta] of Object.entries(result.metafile.outputs)) { if (outputPath.endsWith(".mjs")) { const $outputPath = resolve(workingDir, outputPath); @@ -254,6 +270,8 @@ export async function getBundleResultFromBuild( indexControllerEntryPoint = $outputPath; } else if (isIndexWorkerForTarget(outputMeta.entryPoint, target)) { indexWorkerEntryPoint = $outputPath; + } else if (isInitEntryPoint(outputMeta.entryPoint)) { + initEntryPoint = $outputPath; } else { if ( !outputMeta.entryPoint.startsWith("..") && @@ -280,6 +298,7 @@ export async function getBundleResultFromBuild( runControllerEntryPoint, indexWorkerEntryPoint, indexControllerEntryPoint, + initEntryPoint, contentHash: hasher.digest("hex"), }; } @@ -357,6 +376,7 @@ export async function createBuildManifestFromBundle({ runControllerEntryPoint: bundle.runControllerEntryPoint ?? getRunControllerForTarget(target), runWorkerEntryPoint: bundle.runWorkerEntryPoint ?? getRunWorkerForTarget(target), loaderEntryPoint: bundle.loaderEntryPoint, + initEntryPoint: bundle.initEntryPoint, configPath: bundle.configPath, customConditions: resolvedConfig.build.conditions ?? [], deploy: { diff --git a/packages/cli-v3/src/entryPoints/dev-index-worker.ts b/packages/cli-v3/src/entryPoints/dev-index-worker.ts index a29f2f8541..86528a93fe 100644 --- a/packages/cli-v3/src/entryPoints/dev-index-worker.ts +++ b/packages/cli-v3/src/entryPoints/dev-index-worker.ts @@ -141,6 +141,7 @@ await sendMessageInCatalog( controllerEntryPoint: buildManifest.runControllerEntryPoint, loaderEntryPoint: buildManifest.loaderEntryPoint, customConditions: buildManifest.customConditions, + initEntryPoint: buildManifest.initEntryPoint, }, importErrors, }, diff --git a/packages/cli-v3/src/entryPoints/dev-run-worker.ts b/packages/cli-v3/src/entryPoints/dev-run-worker.ts index 007c001a79..8553e36057 100644 --- a/packages/cli-v3/src/entryPoints/dev-run-worker.ts +++ b/packages/cli-v3/src/entryPoints/dev-run-worker.ts @@ -1,15 +1,23 @@ import type { Tracer } from "@opentelemetry/api"; import type { Logger } from "@opentelemetry/api-logs"; import { + AnyOnCatchErrorHookFunction, + AnyOnFailureHookFunction, + AnyOnInitHookFunction, + AnyOnStartHookFunction, + AnyOnSuccessHookFunction, apiClientManager, clock, ExecutorToWorkerMessageCatalog, type HandleErrorFunction, + lifecycleHooks, + localsAPI, logger, LogLevel, + resourceCatalog, runMetadata, runtime, - resourceCatalog, + runTimelineMetrics, TaskRunErrorCodes, TaskRunExecution, timeout, @@ -17,7 +25,6 @@ import { waitUntil, WorkerManifest, WorkerToExecutorMessageCatalog, - runTimelineMetrics, } from "@trigger.dev/core/v3"; import { TriggerTracer } from "@trigger.dev/core/v3/tracer"; import { @@ -29,15 +36,17 @@ import { logLevels, ManagedRuntimeManager, OtelTaskLogger, + StandardLifecycleHooksManager, + StandardLocalsManager, StandardMetadataManager, StandardResourceCatalog, + StandardRunTimelineMetricsManager, StandardWaitUntilManager, TaskExecutor, TracingDiagnosticLogLevel, TracingSDK, usage, UsageTimeoutManager, - StandardRunTimelineMetricsManager, } from "@trigger.dev/core/v3/workers"; import { ZodIpcConnection } from "@trigger.dev/core/v3/zodIpc"; import { readFile } from "node:fs/promises"; @@ -89,10 +98,16 @@ process.on("uncaughtException", function (error, origin) { const heartbeatIntervalMs = getEnvVar("HEARTBEAT_INTERVAL_MS"); +const standardLocalsManager = new StandardLocalsManager(); +localsAPI.setGlobalLocalsManager(standardLocalsManager); + const standardRunTimelineMetricsManager = new StandardRunTimelineMetricsManager(); runTimelineMetrics.setGlobalManager(standardRunTimelineMetricsManager); standardRunTimelineMetricsManager.seedMetricsFromEnvironment(); +const standardLifecycleHooksManager = new StandardLifecycleHooksManager(); +lifecycleHooks.setGlobalLifecycleHooksManager(standardLifecycleHooksManager); + const devUsageManager = new DevUsageManager(); usage.setGlobalUsageManager(devUsageManager); timeout.setGlobalManager(new UsageTimeoutManager(devUsageManager)); @@ -170,12 +185,46 @@ async function bootstrap() { logger.setGlobalTaskLogger(otelTaskLogger); + if (config.init) { + lifecycleHooks.registerGlobalInitHook({ + id: "config", + fn: config.init as AnyOnInitHookFunction, + }); + } + + if (config.onStart) { + lifecycleHooks.registerGlobalStartHook({ + id: "config", + fn: config.onStart as AnyOnStartHookFunction, + }); + } + + if (config.onSuccess) { + lifecycleHooks.registerGlobalSuccessHook({ + id: "config", + fn: config.onSuccess as AnyOnSuccessHookFunction, + }); + } + + if (config.onFailure) { + lifecycleHooks.registerGlobalFailureHook({ + id: "config", + fn: config.onFailure as AnyOnFailureHookFunction, + }); + } + + if (handleError) { + lifecycleHooks.registerGlobalCatchErrorHook({ + id: "config", + fn: handleError as AnyOnCatchErrorHookFunction, + }); + } + return { tracer, tracingSDK, consoleInterceptor, config, - handleErrorFn: handleError, workerManifest, }; } @@ -217,7 +266,7 @@ const zodIpc = new ZodIpcConnection({ } try { - const { tracer, tracingSDK, consoleInterceptor, config, handleErrorFn, workerManifest } = + const { tracer, tracingSDK, consoleInterceptor, config, workerManifest } = await bootstrap(); _tracingSDK = tracingSDK; @@ -257,6 +306,18 @@ const zodIpc = new ZodIpcConnection({ async () => { const beforeImport = performance.now(); resourceCatalog.setCurrentFileContext(taskManifest.entryPoint, taskManifest.filePath); + + // Load init file if it exists + if (workerManifest.initEntryPoint) { + try { + await import(normalizeImportPath(workerManifest.initEntryPoint)); + log(`Loaded init file from ${workerManifest.initEntryPoint}`); + } catch (err) { + logError(`Failed to load init file`, err); + throw err; + } + } + await import(normalizeImportPath(taskManifest.entryPoint)); resourceCatalog.clearCurrentFileContext(); const durationMs = performance.now() - beforeImport; @@ -321,8 +382,7 @@ const zodIpc = new ZodIpcConnection({ tracer, tracingSDK, consoleInterceptor, - config, - handleErrorFn, + retries: config.retries, }); try { @@ -340,42 +400,7 @@ const zodIpc = new ZodIpcConnection({ ? timeout.abortAfterTimeout(execution.run.maxDuration) : undefined; - signal?.addEventListener("abort", async (e) => { - if (_isRunning) { - _isRunning = false; - _execution = undefined; - - const usageSample = usage.stop(measurement); - - await sender.send("TASK_RUN_COMPLETED", { - execution, - result: { - ok: false, - id: execution.run.id, - error: { - type: "INTERNAL_ERROR", - code: TaskRunErrorCodes.MAX_DURATION_EXCEEDED, - message: - signal.reason instanceof Error - ? signal.reason.message - : String(signal.reason), - }, - usage: { - durationMs: usageSample.cpuTime, - }, - metadata: runMetadataManager.stopAndReturnLastFlush(), - }, - }); - } - }); - - const { result } = await executor.execute( - execution, - metadata, - traceContext, - measurement, - signal - ); + const { result } = await executor.execute(execution, metadata, traceContext, signal); const usageSample = usage.stop(measurement); diff --git a/packages/cli-v3/src/entryPoints/managed-index-worker.ts b/packages/cli-v3/src/entryPoints/managed-index-worker.ts index a29f2f8541..86528a93fe 100644 --- a/packages/cli-v3/src/entryPoints/managed-index-worker.ts +++ b/packages/cli-v3/src/entryPoints/managed-index-worker.ts @@ -141,6 +141,7 @@ await sendMessageInCatalog( controllerEntryPoint: buildManifest.runControllerEntryPoint, loaderEntryPoint: buildManifest.loaderEntryPoint, customConditions: buildManifest.customConditions, + initEntryPoint: buildManifest.initEntryPoint, }, importErrors, }, diff --git a/packages/cli-v3/src/entryPoints/managed-run-worker.ts b/packages/cli-v3/src/entryPoints/managed-run-worker.ts index f261a9e677..fa4f426bac 100644 --- a/packages/cli-v3/src/entryPoints/managed-run-worker.ts +++ b/packages/cli-v3/src/entryPoints/managed-run-worker.ts @@ -1,23 +1,30 @@ import type { Tracer } from "@opentelemetry/api"; import type { Logger } from "@opentelemetry/api-logs"; import { + AnyOnCatchErrorHookFunction, + AnyOnFailureHookFunction, + AnyOnInitHookFunction, + AnyOnStartHookFunction, + AnyOnSuccessHookFunction, + apiClientManager, clock, + ExecutorToWorkerMessageCatalog, type HandleErrorFunction, + lifecycleHooks, + localsAPI, logger, LogLevel, - runtime, resourceCatalog, + runMetadata, + runtime, + runTimelineMetrics, TaskRunErrorCodes, TaskRunExecution, - WorkerToExecutorMessageCatalog, - TriggerConfig, - WorkerManifest, - ExecutorToWorkerMessageCatalog, timeout, - runMetadata, + TriggerConfig, waitUntil, - apiClientManager, - runTimelineMetrics, + WorkerManifest, + WorkerToExecutorMessageCatalog, } from "@trigger.dev/core/v3"; import { TriggerTracer } from "@trigger.dev/core/v3/tracer"; import { @@ -27,18 +34,20 @@ import { getEnvVar, getNumberEnvVar, logLevels, + ManagedRuntimeManager, OtelTaskLogger, ProdUsageManager, + StandardLifecycleHooksManager, + StandardLocalsManager, + StandardMetadataManager, StandardResourceCatalog, + StandardRunTimelineMetricsManager, + StandardWaitUntilManager, TaskExecutor, TracingDiagnosticLogLevel, TracingSDK, usage, UsageTimeoutManager, - StandardMetadataManager, - StandardWaitUntilManager, - ManagedRuntimeManager, - StandardRunTimelineMetricsManager, } from "@trigger.dev/core/v3/workers"; import { ZodIpcConnection } from "@trigger.dev/core/v3/zodIpc"; import { readFile } from "node:fs/promises"; @@ -93,6 +102,12 @@ const usageEventUrl = getEnvVar("USAGE_EVENT_URL"); const triggerJWT = getEnvVar("TRIGGER_JWT"); const heartbeatIntervalMs = getEnvVar("HEARTBEAT_INTERVAL_MS"); +const standardLocalsManager = new StandardLocalsManager(); +localsAPI.setGlobalLocalsManager(standardLocalsManager); + +const standardLifecycleHooksManager = new StandardLifecycleHooksManager(); +lifecycleHooks.setGlobalLifecycleHooksManager(standardLifecycleHooksManager); + const standardRunTimelineMetricsManager = new StandardRunTimelineMetricsManager(); runTimelineMetrics.setGlobalManager(standardRunTimelineMetricsManager); standardRunTimelineMetricsManager.seedMetricsFromEnvironment(); @@ -180,12 +195,46 @@ async function bootstrap() { logger.setGlobalTaskLogger(otelTaskLogger); + if (config.init) { + lifecycleHooks.registerGlobalInitHook({ + id: "config", + fn: config.init as AnyOnInitHookFunction, + }); + } + + if (config.onStart) { + lifecycleHooks.registerGlobalStartHook({ + id: "config", + fn: config.onStart as AnyOnStartHookFunction, + }); + } + + if (config.onSuccess) { + lifecycleHooks.registerGlobalSuccessHook({ + id: "config", + fn: config.onSuccess as AnyOnSuccessHookFunction, + }); + } + + if (config.onFailure) { + lifecycleHooks.registerGlobalFailureHook({ + id: "config", + fn: config.onFailure as AnyOnFailureHookFunction, + }); + } + + if (handleError) { + lifecycleHooks.registerGlobalCatchErrorHook({ + id: "config", + fn: handleError as AnyOnCatchErrorHookFunction, + }); + } + return { tracer, tracingSDK, consoleInterceptor, config, - handleErrorFn: handleError, workerManifest, }; } @@ -227,7 +276,7 @@ const zodIpc = new ZodIpcConnection({ } try { - const { tracer, tracingSDK, consoleInterceptor, config, handleErrorFn, workerManifest } = + const { tracer, tracingSDK, consoleInterceptor, config, workerManifest } = await bootstrap(); _tracingSDK = tracingSDK; @@ -331,8 +380,7 @@ const zodIpc = new ZodIpcConnection({ tracer, tracingSDK, consoleInterceptor, - config, - handleErrorFn, + retries: config.retries, }); try { @@ -350,42 +398,7 @@ const zodIpc = new ZodIpcConnection({ ? timeout.abortAfterTimeout(execution.run.maxDuration) : undefined; - signal?.addEventListener("abort", async (e) => { - if (_isRunning) { - _isRunning = false; - _execution = undefined; - - const usageSample = usage.stop(measurement); - - await sender.send("TASK_RUN_COMPLETED", { - execution, - result: { - ok: false, - id: execution.run.id, - error: { - type: "INTERNAL_ERROR", - code: TaskRunErrorCodes.MAX_DURATION_EXCEEDED, - message: - signal.reason instanceof Error - ? signal.reason.message - : String(signal.reason), - }, - usage: { - durationMs: usageSample.cpuTime, - }, - metadata: runMetadataManager.stopAndReturnLastFlush(), - }, - }); - } - }); - - const { result } = await executor.execute( - execution, - metadata, - traceContext, - measurement, - signal - ); + const { result } = await executor.execute(execution, metadata, traceContext, signal); const usageSample = usage.stop(measurement); diff --git a/packages/core/src/v3/config.ts b/packages/core/src/v3/config.ts index 3862a24c90..9be80decc6 100644 --- a/packages/core/src/v3/config.ts +++ b/packages/core/src/v3/config.ts @@ -1,15 +1,16 @@ import type { Instrumentation } from "@opentelemetry/instrumentation"; import type { SpanExporter } from "@opentelemetry/sdk-trace-base"; import type { BuildExtension } from "./build/extensions.js"; -import type { MachinePresetName } from "./schemas/common.js"; -import type { LogLevel } from "./logger/taskLogger.js"; import type { - FailureFnParams, - InitFnParams, - StartFnParams, - SuccessFnParams, -} from "./types/index.js"; -import type { BuildRuntime, RetryOptions } from "./index.js"; + AnyOnFailureHookFunction, + AnyOnInitHookFunction, + AnyOnStartHookFunction, + AnyOnSuccessHookFunction, + BuildRuntime, + RetryOptions, +} from "./index.js"; +import type { LogLevel } from "./logger/taskLogger.js"; +import type { MachinePresetName } from "./schemas/common.js"; export type CompatibilityFlag = "run_engine_v2"; @@ -215,23 +216,31 @@ export type TriggerConfig = { /** * Run before a task is executed, for all tasks. This is useful for setting up any global state that is needed for all tasks. + * + * @deprecated, please use tasks.init instead */ - init?: (payload: unknown, params: InitFnParams) => void | Promise; + init?: AnyOnInitHookFunction; /** * onSuccess is called after the run function has successfully completed. + * + * @deprecated, please use tasks.onSuccess instead */ - onSuccess?: (payload: unknown, output: unknown, params: SuccessFnParams) => Promise; + onSuccess?: AnyOnSuccessHookFunction; /** * onFailure is called after a task run has failed (meaning the run function threw an error and won't be retried anymore) + * + * @deprecated, please use tasks.onFailure instead */ - onFailure?: (payload: unknown, error: unknown, params: FailureFnParams) => Promise; + onFailure?: AnyOnFailureHookFunction; /** * onStart is called the first time a task is executed in a run (not before every retry) + * + * @deprecated, please use tasks.onStart instead */ - onStart?: (payload: unknown, params: StartFnParams) => Promise; + onStart?: AnyOnStartHookFunction; /** * @deprecated Use a custom build extension to add post install commands diff --git a/packages/core/src/v3/errors.ts b/packages/core/src/v3/errors.ts index a8e789f81d..bf9776d1c2 100644 --- a/packages/core/src/v3/errors.ts +++ b/packages/core/src/v3/errors.ts @@ -305,6 +305,7 @@ export function shouldRetryError(error: TaskRunError): boolean { case "HANDLE_ERROR_ERROR": case "TASK_INPUT_ERROR": case "TASK_OUTPUT_ERROR": + case "TASK_MIDDLEWARE_ERROR": case "POD_EVICTED": case "POD_UNKNOWN_ERROR": case "TASK_EXECUTION_ABORTED": diff --git a/packages/core/src/v3/index.ts b/packages/core/src/v3/index.ts index 098362a7dc..7706842c14 100644 --- a/packages/core/src/v3/index.ts +++ b/packages/core/src/v3/index.ts @@ -15,6 +15,8 @@ export * from "./run-metadata-api.js"; export * from "./wait-until-api.js"; export * from "./timeout-api.js"; export * from "./run-timeline-metrics-api.js"; +export * from "./lifecycle-hooks-api.js"; +export * from "./locals-api.js"; export * from "./schemas/index.js"; export { SemanticInternalAttributes } from "./semanticInternalAttributes.js"; export * from "./resource-catalog-api.js"; diff --git a/packages/core/src/v3/lifecycle-hooks-api.ts b/packages/core/src/v3/lifecycle-hooks-api.ts new file mode 100644 index 0000000000..ec9e87c998 --- /dev/null +++ b/packages/core/src/v3/lifecycle-hooks-api.ts @@ -0,0 +1,35 @@ +// Split module-level variable definition into separate files to allow +// tree-shaking on each api instance. +import { LifecycleHooksAPI } from "./lifecycleHooks/index.js"; +/** Entrypoint for runtime API */ +export const lifecycleHooks = LifecycleHooksAPI.getInstance(); + +export type { + OnInitHookFunction, + AnyOnInitHookFunction, + RegisteredHookFunction, + TaskInitHookParams, + TaskStartHookParams, + OnStartHookFunction, + AnyOnStartHookFunction, + TaskFailureHookParams, + AnyOnFailureHookFunction, + TaskSuccessHookParams, + AnyOnSuccessHookFunction, + TaskCompleteHookParams, + AnyOnCompleteHookFunction, + TaskWaitHookParams, + AnyOnWaitHookFunction, + TaskResumeHookParams, + AnyOnResumeHookFunction, + TaskCatchErrorHookParams, + AnyOnCatchErrorHookFunction, + TaskCompleteResult, + TaskMiddlewareHookParams, + AnyOnMiddlewareHookFunction, + OnMiddlewareHookFunction, + OnCleanupHookFunction, + AnyOnCleanupHookFunction, + TaskCleanupHookParams, + TaskWait, +} from "./lifecycleHooks/types.js"; diff --git a/packages/core/src/v3/lifecycleHooks/index.ts b/packages/core/src/v3/lifecycleHooks/index.ts new file mode 100644 index 0000000000..843ae92ce8 --- /dev/null +++ b/packages/core/src/v3/lifecycleHooks/index.ts @@ -0,0 +1,266 @@ +const API_NAME = "lifecycle-hooks"; + +import { getGlobal, registerGlobal, unregisterGlobal } from "../utils/globals.js"; +import { NoopLifecycleHooksManager } from "./manager.js"; +import { + AnyOnCatchErrorHookFunction, + AnyOnCleanupHookFunction, + AnyOnCompleteHookFunction, + AnyOnFailureHookFunction, + AnyOnInitHookFunction, + AnyOnMiddlewareHookFunction, + AnyOnResumeHookFunction, + AnyOnStartHookFunction, + AnyOnSuccessHookFunction, + AnyOnWaitHookFunction, + RegisteredHookFunction, + RegisterHookFunctionParams, + TaskWait, + type LifecycleHooksManager, +} from "./types.js"; + +const NOOP_LIFECYCLE_HOOKS_MANAGER = new NoopLifecycleHooksManager(); + +export class LifecycleHooksAPI { + private static _instance?: LifecycleHooksAPI; + + private constructor() {} + + public static getInstance(): LifecycleHooksAPI { + if (!this._instance) { + this._instance = new LifecycleHooksAPI(); + } + + return this._instance; + } + + public setGlobalLifecycleHooksManager(lifecycleHooksManager: LifecycleHooksManager): boolean { + return registerGlobal(API_NAME, lifecycleHooksManager); + } + + public disable() { + unregisterGlobal(API_NAME); + } + + public registerGlobalInitHook(hook: RegisterHookFunctionParams): void { + this.#getManager().registerGlobalInitHook(hook); + } + + public registerTaskInitHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerTaskInitHook(taskId, hook); + } + + public getTaskInitHook(taskId: string): AnyOnInitHookFunction | undefined { + return this.#getManager().getTaskInitHook(taskId); + } + + public getGlobalInitHooks(): RegisteredHookFunction[] { + return this.#getManager().getGlobalInitHooks(); + } + + public registerTaskStartHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerTaskStartHook(taskId, hook); + } + + public registerGlobalStartHook(hook: RegisterHookFunctionParams): void { + this.#getManager().registerGlobalStartHook(hook); + } + + public getTaskStartHook(taskId: string): AnyOnStartHookFunction | undefined { + return this.#getManager().getTaskStartHook(taskId); + } + + public getGlobalStartHooks(): RegisteredHookFunction[] { + return this.#getManager().getGlobalStartHooks(); + } + + public registerGlobalFailureHook( + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerGlobalFailureHook(hook); + } + + public registerTaskFailureHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerTaskFailureHook(taskId, hook); + } + + public getTaskFailureHook(taskId: string): AnyOnFailureHookFunction | undefined { + return this.#getManager().getTaskFailureHook(taskId); + } + + public getGlobalFailureHooks(): RegisteredHookFunction[] { + return this.#getManager().getGlobalFailureHooks(); + } + + public registerGlobalSuccessHook( + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerGlobalSuccessHook(hook); + } + + public registerTaskSuccessHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerTaskSuccessHook(taskId, hook); + } + + public getTaskSuccessHook(taskId: string): AnyOnSuccessHookFunction | undefined { + return this.#getManager().getTaskSuccessHook(taskId); + } + + public getGlobalSuccessHooks(): RegisteredHookFunction[] { + return this.#getManager().getGlobalSuccessHooks(); + } + + public registerGlobalCompleteHook( + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerGlobalCompleteHook(hook); + } + + public registerTaskCompleteHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerTaskCompleteHook(taskId, hook); + } + + public getTaskCompleteHook(taskId: string): AnyOnCompleteHookFunction | undefined { + return this.#getManager().getTaskCompleteHook(taskId); + } + + public getGlobalCompleteHooks(): RegisteredHookFunction[] { + return this.#getManager().getGlobalCompleteHooks(); + } + + public registerGlobalWaitHook(hook: RegisterHookFunctionParams): void { + this.#getManager().registerGlobalWaitHook(hook); + } + + public registerTaskWaitHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerTaskWaitHook(taskId, hook); + } + + public getTaskWaitHook(taskId: string): AnyOnWaitHookFunction | undefined { + return this.#getManager().getTaskWaitHook(taskId); + } + + public getGlobalWaitHooks(): RegisteredHookFunction[] { + return this.#getManager().getGlobalWaitHooks(); + } + + public registerGlobalResumeHook(hook: RegisterHookFunctionParams): void { + this.#getManager().registerGlobalResumeHook(hook); + } + + public registerTaskResumeHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerTaskResumeHook(taskId, hook); + } + + public getTaskResumeHook(taskId: string): AnyOnResumeHookFunction | undefined { + return this.#getManager().getTaskResumeHook(taskId); + } + + public getGlobalResumeHooks(): RegisteredHookFunction[] { + return this.#getManager().getGlobalResumeHooks(); + } + + public registerGlobalCatchErrorHook( + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerGlobalCatchErrorHook(hook); + } + + public registerTaskCatchErrorHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerTaskCatchErrorHook(taskId, hook); + } + + public getTaskCatchErrorHook(taskId: string): AnyOnCatchErrorHookFunction | undefined { + return this.#getManager().getTaskCatchErrorHook(taskId); + } + + public getGlobalCatchErrorHooks(): RegisteredHookFunction[] { + return this.#getManager().getGlobalCatchErrorHooks(); + } + + public registerGlobalMiddlewareHook( + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerGlobalMiddlewareHook(hook); + } + + public registerTaskMiddlewareHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerTaskMiddlewareHook(taskId, hook); + } + + public getTaskMiddlewareHook(taskId: string): AnyOnMiddlewareHookFunction | undefined { + return this.#getManager().getTaskMiddlewareHook(taskId); + } + + public getGlobalMiddlewareHooks(): RegisteredHookFunction[] { + return this.#getManager().getGlobalMiddlewareHooks(); + } + + public registerGlobalCleanupHook( + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerGlobalCleanupHook(hook); + } + + public registerTaskCleanupHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + this.#getManager().registerTaskCleanupHook(taskId, hook); + } + + public getTaskCleanupHook(taskId: string): AnyOnCleanupHookFunction | undefined { + return this.#getManager().getTaskCleanupHook(taskId); + } + + public getGlobalCleanupHooks(): RegisteredHookFunction[] { + return this.#getManager().getGlobalCleanupHooks(); + } + + public callOnWaitHookListeners(wait: TaskWait): Promise { + return this.#getManager().callOnWaitHookListeners(wait); + } + + public callOnResumeHookListeners(wait: TaskWait): Promise { + return this.#getManager().callOnResumeHookListeners(wait); + } + + public registerOnWaitHookListener(listener: (wait: TaskWait) => Promise): void { + this.#getManager().registerOnWaitHookListener(listener); + } + + public registerOnResumeHookListener(listener: (wait: TaskWait) => Promise): void { + this.#getManager().registerOnResumeHookListener(listener); + } + + #getManager(): LifecycleHooksManager { + return getGlobal(API_NAME) ?? NOOP_LIFECYCLE_HOOKS_MANAGER; + } +} diff --git a/packages/core/src/v3/lifecycleHooks/manager.ts b/packages/core/src/v3/lifecycleHooks/manager.ts new file mode 100644 index 0000000000..29f4968362 --- /dev/null +++ b/packages/core/src/v3/lifecycleHooks/manager.ts @@ -0,0 +1,603 @@ +import { + AnyOnInitHookFunction, + AnyOnStartHookFunction, + LifecycleHooksManager, + RegisteredHookFunction, + RegisterHookFunctionParams, + AnyOnFailureHookFunction, + AnyOnSuccessHookFunction, + AnyOnCompleteHookFunction, + AnyOnWaitHookFunction, + AnyOnResumeHookFunction, + AnyOnCatchErrorHookFunction, + AnyOnMiddlewareHookFunction, + AnyOnCleanupHookFunction, + TaskWait, +} from "./types.js"; + +export class StandardLifecycleHooksManager implements LifecycleHooksManager { + private globalInitHooks: Map> = new Map(); + private taskInitHooks: Map> = new Map(); + + private globalStartHooks: Map> = new Map(); + private taskStartHooks: Map> = new Map(); + + private globalFailureHooks: Map> = + new Map(); + private taskFailureHooks: Map> = + new Map(); + + private globalSuccessHooks: Map> = + new Map(); + private taskSuccessHooks: Map> = + new Map(); + + private globalCompleteHooks: Map> = + new Map(); + private taskCompleteHooks: Map> = + new Map(); + + private globalWaitHooks: Map> = new Map(); + private taskWaitHooks: Map> = new Map(); + + private globalResumeHooks: Map> = + new Map(); + private taskResumeHooks: Map> = new Map(); + + private globalCatchErrorHooks: Map> = + new Map(); + private taskCatchErrorHooks: Map> = + new Map(); + + private globalMiddlewareHooks: Map> = + new Map(); + private taskMiddlewareHooks: Map> = + new Map(); + + private globalCleanupHooks: Map> = + new Map(); + private taskCleanupHooks: Map> = + new Map(); + + private onWaitHookListeners: ((wait: TaskWait) => Promise)[] = []; + private onResumeHookListeners: ((wait: TaskWait) => Promise)[] = []; + + registerOnWaitHookListener(listener: (wait: TaskWait) => Promise): void { + this.onWaitHookListeners.push(listener); + } + + async callOnWaitHookListeners(wait: TaskWait): Promise { + await Promise.allSettled(this.onWaitHookListeners.map((listener) => listener(wait))); + } + + registerOnResumeHookListener(listener: (wait: TaskWait) => Promise): void { + this.onResumeHookListeners.push(listener); + } + + async callOnResumeHookListeners(wait: TaskWait): Promise { + await Promise.allSettled(this.onResumeHookListeners.map((listener) => listener(wait))); + } + + registerGlobalStartHook(hook: RegisterHookFunctionParams): void { + const id = generateHookId(hook); + + this.globalStartHooks.set(id, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + registerTaskStartHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.taskStartHooks.set(taskId, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + getTaskStartHook(taskId: string): AnyOnStartHookFunction | undefined { + return this.taskStartHooks.get(taskId)?.fn; + } + + getGlobalStartHooks(): RegisteredHookFunction[] { + return Array.from(this.globalStartHooks.values()); + } + + registerGlobalInitHook(hook: RegisterHookFunctionParams): void { + // if there is no id, lets generate one based on the contents of the function + const id = generateHookId(hook); + + const registeredHook = { + id, + name: hook.id, + fn: hook.fn, + }; + + this.globalInitHooks.set(id, registeredHook); + } + + registerTaskInitHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const registeredHook = { + id: generateHookId(hook), + name: taskId, + fn: hook.fn, + }; + + this.taskInitHooks.set(taskId, registeredHook); + } + + getTaskInitHook(taskId: string): AnyOnInitHookFunction | undefined { + return this.taskInitHooks.get(taskId)?.fn; + } + + getGlobalInitHooks(): RegisteredHookFunction[] { + return Array.from(this.globalInitHooks.values()); + } + + registerGlobalFailureHook(hook: RegisterHookFunctionParams): void { + const id = generateHookId(hook); + + this.globalFailureHooks.set(id, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + registerTaskFailureHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.taskFailureHooks.set(taskId, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + getTaskFailureHook(taskId: string): AnyOnFailureHookFunction | undefined { + return this.taskFailureHooks.get(taskId)?.fn; + } + + getGlobalFailureHooks(): RegisteredHookFunction[] { + return Array.from(this.globalFailureHooks.values()); + } + + registerGlobalSuccessHook(hook: RegisterHookFunctionParams): void { + const id = generateHookId(hook); + + this.globalSuccessHooks.set(id, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + registerTaskSuccessHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.taskSuccessHooks.set(taskId, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + getTaskSuccessHook(taskId: string): AnyOnSuccessHookFunction | undefined { + return this.taskSuccessHooks.get(taskId)?.fn; + } + + getGlobalSuccessHooks(): RegisteredHookFunction[] { + return Array.from(this.globalSuccessHooks.values()); + } + + registerGlobalCompleteHook(hook: RegisterHookFunctionParams): void { + const id = generateHookId(hook); + + this.globalCompleteHooks.set(id, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + registerTaskCompleteHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.taskCompleteHooks.set(taskId, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + getTaskCompleteHook(taskId: string): AnyOnCompleteHookFunction | undefined { + return this.taskCompleteHooks.get(taskId)?.fn; + } + + getGlobalCompleteHooks(): RegisteredHookFunction[] { + return Array.from(this.globalCompleteHooks.values()); + } + + registerGlobalWaitHook(hook: RegisterHookFunctionParams): void { + const id = generateHookId(hook); + + this.globalWaitHooks.set(id, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + registerTaskWaitHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.taskWaitHooks.set(taskId, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + getTaskWaitHook(taskId: string): AnyOnWaitHookFunction | undefined { + return this.taskWaitHooks.get(taskId)?.fn; + } + + getGlobalWaitHooks(): RegisteredHookFunction[] { + return Array.from(this.globalWaitHooks.values()); + } + + registerGlobalResumeHook(hook: RegisterHookFunctionParams): void { + const id = generateHookId(hook); + + this.globalResumeHooks.set(id, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + registerTaskResumeHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.taskResumeHooks.set(taskId, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + getTaskResumeHook(taskId: string): AnyOnResumeHookFunction | undefined { + return this.taskResumeHooks.get(taskId)?.fn; + } + + getGlobalResumeHooks(): RegisteredHookFunction[] { + return Array.from(this.globalResumeHooks.values()); + } + + registerGlobalCatchErrorHook( + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.globalCatchErrorHooks.set(id, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + registerTaskCatchErrorHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.taskCatchErrorHooks.set(taskId, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + getTaskCatchErrorHook(taskId: string): AnyOnCatchErrorHookFunction | undefined { + return this.taskCatchErrorHooks.get(taskId)?.fn; + } + + getGlobalCatchErrorHooks(): RegisteredHookFunction[] { + return Array.from(this.globalCatchErrorHooks.values()); + } + + registerGlobalMiddlewareHook( + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.globalMiddlewareHooks.set(id, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + registerTaskMiddlewareHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.taskMiddlewareHooks.set(taskId, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + getTaskMiddlewareHook(taskId: string): AnyOnMiddlewareHookFunction | undefined { + return this.taskMiddlewareHooks.get(taskId)?.fn; + } + + getGlobalMiddlewareHooks(): RegisteredHookFunction[] { + return Array.from(this.globalMiddlewareHooks.values()); + } + + registerGlobalCleanupHook(hook: RegisterHookFunctionParams): void { + const id = generateHookId(hook); + + this.globalCleanupHooks.set(id, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + registerTaskCleanupHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + const id = generateHookId(hook); + + this.taskCleanupHooks.set(taskId, { + id, + name: hook.id, + fn: hook.fn, + }); + } + + getTaskCleanupHook(taskId: string): AnyOnCleanupHookFunction | undefined { + return this.taskCleanupHooks.get(taskId)?.fn; + } + + getGlobalCleanupHooks(): RegisteredHookFunction[] { + return Array.from(this.globalCleanupHooks.values()); + } +} + +export class NoopLifecycleHooksManager implements LifecycleHooksManager { + registerOnWaitHookListener(listener: (wait: TaskWait) => Promise): void { + // Noop + } + + async callOnWaitHookListeners(wait: TaskWait): Promise { + // Noop + } + + registerOnResumeHookListener(listener: (wait: TaskWait) => Promise): void { + // Noop + } + + async callOnResumeHookListeners(wait: TaskWait): Promise { + // Noop + } + + registerGlobalInitHook(hook: RegisterHookFunctionParams): void { + // Noop + } + + registerTaskInitHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + // Noop + } + + getTaskInitHook(taskId: string): AnyOnInitHookFunction | undefined { + return undefined; + } + + getGlobalInitHooks(): RegisteredHookFunction[] { + return []; + } + + registerGlobalStartHook(hook: RegisterHookFunctionParams): void { + // Noop + } + + registerTaskStartHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + // Noop + } + + getTaskStartHook(taskId: string): AnyOnStartHookFunction | undefined { + return undefined; + } + + getGlobalStartHooks(): RegisteredHookFunction[] { + return []; + } + + registerGlobalFailureHook(hook: RegisterHookFunctionParams): void { + // Noop + } + + registerTaskFailureHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + // Noop + } + + getTaskFailureHook(taskId: string): AnyOnFailureHookFunction | undefined { + return undefined; + } + + getGlobalFailureHooks(): RegisteredHookFunction[] { + return []; + } + + registerGlobalSuccessHook(hook: RegisterHookFunctionParams): void { + // Noop + } + + registerTaskSuccessHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + // Noop + } + + getTaskSuccessHook(taskId: string): AnyOnSuccessHookFunction | undefined { + return undefined; + } + + getGlobalSuccessHooks(): RegisteredHookFunction[] { + return []; + } + + registerGlobalCompleteHook(hook: RegisterHookFunctionParams): void { + // Noop + } + + registerTaskCompleteHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + // Noop + } + + getTaskCompleteHook(taskId: string): AnyOnCompleteHookFunction | undefined { + return undefined; + } + + getGlobalCompleteHooks(): RegisteredHookFunction[] { + return []; + } + + registerGlobalWaitHook(hook: RegisterHookFunctionParams): void { + // Noop + } + + registerTaskWaitHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + // Noop + } + + getTaskWaitHook(taskId: string): AnyOnWaitHookFunction | undefined { + return undefined; + } + + getGlobalWaitHooks(): RegisteredHookFunction[] { + return []; + } + + registerGlobalResumeHook(hook: RegisterHookFunctionParams): void { + // Noop + } + + registerTaskResumeHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + // Noop + } + + getTaskResumeHook(taskId: string): AnyOnResumeHookFunction | undefined { + return undefined; + } + + getGlobalResumeHooks(): RegisteredHookFunction[] { + return []; + } + + registerGlobalCatchErrorHook(): void { + // Noop + } + + registerTaskCatchErrorHook(): void { + // Noop + } + + getTaskCatchErrorHook(): undefined { + return undefined; + } + + getGlobalCatchErrorHooks(): [] { + return []; + } + + registerGlobalMiddlewareHook(): void { + // Noop + } + + registerTaskMiddlewareHook(): void { + // Noop + } + + getTaskMiddlewareHook(): undefined { + return undefined; + } + + getGlobalMiddlewareHooks(): [] { + return []; + } + + registerGlobalCleanupHook(hook: RegisterHookFunctionParams): void { + // Noop + } + + registerTaskCleanupHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void { + // Noop + } + + getTaskCleanupHook(taskId: string): AnyOnCleanupHookFunction | undefined { + return undefined; + } + + getGlobalCleanupHooks(): RegisteredHookFunction[] { + return []; + } +} + +function generateHookId(hook: RegisterHookFunctionParams): string { + return hook.id ?? hook.fn.toString(); +} diff --git a/packages/core/src/v3/lifecycleHooks/types.ts b/packages/core/src/v3/lifecycleHooks/types.ts new file mode 100644 index 0000000000..5d307c225b --- /dev/null +++ b/packages/core/src/v3/lifecycleHooks/types.ts @@ -0,0 +1,310 @@ +import { RetryOptions, TaskRunContext } from "../schemas/index.js"; +import { HandleErrorResult } from "../types/index.js"; + +export type TaskInitOutput = Record | void | undefined; + +export type TaskInitHookParams = { + ctx: TaskRunContext; + payload: TPayload; + task: string; + signal?: AbortSignal; +}; + +export type OnInitHookFunction = ( + params: TaskInitHookParams +) => TInitOutput | undefined | void | Promise; + +export type AnyOnInitHookFunction = OnInitHookFunction; + +export type TaskStartHookParams< + TPayload = unknown, + TInitOutput extends TaskInitOutput = TaskInitOutput, +> = { + ctx: TaskRunContext; + payload: TPayload; + task: string; + signal?: AbortSignal; + init?: TInitOutput; +}; + +export type OnStartHookFunction = ( + params: TaskStartHookParams +) => undefined | void | Promise; + +export type AnyOnStartHookFunction = OnStartHookFunction; + +export type TaskWait = + | { + type: "duration"; + date: Date; + } + | { + type: "token"; + token: string; + } + | { + type: "task"; + runId: string; + } + | { + type: "batch"; + batchId: string; + runCount: number; + }; + +export type TaskWaitHookParams< + TPayload = unknown, + TInitOutput extends TaskInitOutput = TaskInitOutput, +> = { + wait: TaskWait; + ctx: TaskRunContext; + payload: TPayload; + task: string; + signal?: AbortSignal; + init?: TInitOutput; +}; + +export type OnWaitHookFunction = ( + params: TaskWaitHookParams +) => undefined | void | Promise; + +export type AnyOnWaitHookFunction = OnWaitHookFunction; + +export type TaskResumeHookParams< + TPayload = unknown, + TInitOutput extends TaskInitOutput = TaskInitOutput, +> = { + ctx: TaskRunContext; + wait: TaskWait; + payload: TPayload; + task: string; + signal?: AbortSignal; + init?: TInitOutput; +}; + +export type OnResumeHookFunction = ( + params: TaskResumeHookParams +) => undefined | void | Promise; + +export type AnyOnResumeHookFunction = OnResumeHookFunction; + +export type TaskFailureHookParams< + TPayload = unknown, + TInitOutput extends TaskInitOutput = TaskInitOutput, +> = { + ctx: TaskRunContext; + payload: TPayload; + task: string; + error: unknown; + signal?: AbortSignal; + init?: TInitOutput; +}; + +export type OnFailureHookFunction = ( + params: TaskFailureHookParams +) => undefined | void | Promise; + +export type AnyOnFailureHookFunction = OnFailureHookFunction; + +export type TaskSuccessHookParams< + TPayload = unknown, + TOutput = unknown, + TInitOutput extends TaskInitOutput = TaskInitOutput, +> = { + ctx: TaskRunContext; + payload: TPayload; + task: string; + output: TOutput; + signal?: AbortSignal; + init?: TInitOutput; +}; + +export type OnSuccessHookFunction< + TPayload, + TOutput, + TInitOutput extends TaskInitOutput = TaskInitOutput, +> = ( + params: TaskSuccessHookParams +) => undefined | void | Promise; + +export type AnyOnSuccessHookFunction = OnSuccessHookFunction; + +export type TaskCompleteSuccessResult = { + ok: true; + data: TOutput; +}; + +export type TaskCompleteErrorResult = { + ok: false; + error: unknown; +}; + +export type TaskCompleteResult = + | TaskCompleteSuccessResult + | TaskCompleteErrorResult; + +export type TaskCompleteHookParams< + TPayload = unknown, + TOutput = unknown, + TInitOutput extends TaskInitOutput = TaskInitOutput, +> = { + ctx: TaskRunContext; + payload: TPayload; + task: string; + result: TaskCompleteResult; + signal?: AbortSignal; + init?: TInitOutput; +}; + +export type OnCompleteHookFunction< + TPayload, + TOutput, + TInitOutput extends TaskInitOutput = TaskInitOutput, +> = ( + params: TaskCompleteHookParams +) => undefined | void | Promise; + +export type AnyOnCompleteHookFunction = OnCompleteHookFunction; + +export type RegisterHookFunctionParams any> = { + id?: string; + fn: THookFunction; +}; + +export type RegisteredHookFunction any> = { + id: string; + name?: string; + fn: THookFunction; +}; + +export type TaskCatchErrorHookParams< + TPayload = unknown, + TInitOutput extends TaskInitOutput = TaskInitOutput, +> = { + ctx: TaskRunContext; + payload: TPayload; + task: string; + error: unknown; + retry?: RetryOptions; + retryAt?: Date; + retryDelayInMs?: number; + signal?: AbortSignal; + init?: TInitOutput; +}; + +export type OnCatchErrorHookFunction< + TPayload, + TInitOutput extends TaskInitOutput = TaskInitOutput, +> = (params: TaskCatchErrorHookParams) => HandleErrorResult; + +export type AnyOnCatchErrorHookFunction = OnCatchErrorHookFunction; + +export type TaskMiddlewareHookParams = { + ctx: TaskRunContext; + payload: TPayload; + task: string; + signal?: AbortSignal; + next: () => Promise; +}; + +export type OnMiddlewareHookFunction = ( + params: TaskMiddlewareHookParams +) => Promise; + +export type AnyOnMiddlewareHookFunction = OnMiddlewareHookFunction; + +export type TaskCleanupHookParams< + TPayload = unknown, + TInitOutput extends TaskInitOutput = TaskInitOutput, +> = { + ctx: TaskRunContext; + payload: TPayload; + task: string; + signal?: AbortSignal; + init?: TInitOutput; +}; + +export type OnCleanupHookFunction = ( + params: TaskCleanupHookParams +) => undefined | void | Promise; + +export type AnyOnCleanupHookFunction = OnCleanupHookFunction; + +export interface LifecycleHooksManager { + registerGlobalInitHook(hook: RegisterHookFunctionParams): void; + registerTaskInitHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskInitHook(taskId: string): AnyOnInitHookFunction | undefined; + getGlobalInitHooks(): RegisteredHookFunction[]; + registerGlobalStartHook(hook: RegisterHookFunctionParams): void; + registerTaskStartHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskStartHook(taskId: string): AnyOnStartHookFunction | undefined; + getGlobalStartHooks(): RegisteredHookFunction[]; + registerGlobalFailureHook(hook: RegisterHookFunctionParams): void; + registerTaskFailureHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskFailureHook(taskId: string): AnyOnFailureHookFunction | undefined; + getGlobalFailureHooks(): RegisteredHookFunction[]; + registerGlobalSuccessHook(hook: RegisterHookFunctionParams): void; + registerTaskSuccessHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskSuccessHook(taskId: string): AnyOnSuccessHookFunction | undefined; + getGlobalSuccessHooks(): RegisteredHookFunction[]; + registerGlobalCompleteHook(hook: RegisterHookFunctionParams): void; + registerTaskCompleteHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskCompleteHook(taskId: string): AnyOnCompleteHookFunction | undefined; + getGlobalCompleteHooks(): RegisteredHookFunction[]; + registerGlobalWaitHook(hook: RegisterHookFunctionParams): void; + registerTaskWaitHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskWaitHook(taskId: string): AnyOnWaitHookFunction | undefined; + getGlobalWaitHooks(): RegisteredHookFunction[]; + registerGlobalResumeHook(hook: RegisterHookFunctionParams): void; + registerTaskResumeHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskResumeHook(taskId: string): AnyOnResumeHookFunction | undefined; + getGlobalResumeHooks(): RegisteredHookFunction[]; + registerGlobalCatchErrorHook(hook: RegisterHookFunctionParams): void; + registerTaskCatchErrorHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskCatchErrorHook(taskId: string): AnyOnCatchErrorHookFunction | undefined; + getGlobalCatchErrorHooks(): RegisteredHookFunction[]; + registerGlobalMiddlewareHook(hook: RegisterHookFunctionParams): void; + registerTaskMiddlewareHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskMiddlewareHook(taskId: string): AnyOnMiddlewareHookFunction | undefined; + getGlobalMiddlewareHooks(): RegisteredHookFunction[]; + registerGlobalCleanupHook(hook: RegisterHookFunctionParams): void; + registerTaskCleanupHook( + taskId: string, + hook: RegisterHookFunctionParams + ): void; + getTaskCleanupHook(taskId: string): AnyOnCleanupHookFunction | undefined; + getGlobalCleanupHooks(): RegisteredHookFunction[]; + + callOnWaitHookListeners(wait: TaskWait): Promise; + registerOnWaitHookListener(listener: (wait: TaskWait) => Promise): void; + + callOnResumeHookListeners(wait: TaskWait): Promise; + registerOnResumeHookListener(listener: (wait: TaskWait) => Promise): void; +} diff --git a/packages/core/src/v3/locals-api.ts b/packages/core/src/v3/locals-api.ts new file mode 100644 index 0000000000..a9f86494a6 --- /dev/null +++ b/packages/core/src/v3/locals-api.ts @@ -0,0 +1,29 @@ +// Split module-level variable definition into separate files to allow +// tree-shaking on each api instance. +import { LocalsAPI } from "./locals/index.js"; +import type { LocalsKey } from "./locals/types.js"; +/** Entrypoint for runtime API */ +export const localsAPI = LocalsAPI.getInstance(); + +export const locals = { + create(id: string): LocalsKey { + return localsAPI.createLocal(id); + }, + get(key: LocalsKey): T | undefined { + return localsAPI.getLocal(key); + }, + getOrThrow(key: LocalsKey): T { + const value = localsAPI.getLocal(key); + if (!value) { + throw new Error(`Local with id ${key.id} not found`); + } + return value; + }, + set(key: LocalsKey, value: T): T { + localsAPI.setLocal(key, value); + return value; + }, +}; + +export type Locals = typeof locals; +export type { LocalsKey }; diff --git a/packages/core/src/v3/locals/index.ts b/packages/core/src/v3/locals/index.ts new file mode 100644 index 0000000000..def8602384 --- /dev/null +++ b/packages/core/src/v3/locals/index.ts @@ -0,0 +1,45 @@ +const API_NAME = "locals"; + +import { getGlobal, registerGlobal, unregisterGlobal } from "../utils/globals.js"; +import { NoopLocalsManager } from "./manager.js"; +import { LocalsKey, type LocalsManager } from "./types.js"; + +const NOOP_LOCALS_MANAGER = new NoopLocalsManager(); + +export class LocalsAPI implements LocalsManager { + private static _instance?: LocalsAPI; + + private constructor() {} + + public static getInstance(): LocalsAPI { + if (!this._instance) { + this._instance = new LocalsAPI(); + } + + return this._instance; + } + + public setGlobalLocalsManager(localsManager: LocalsManager): boolean { + return registerGlobal(API_NAME, localsManager); + } + + public disable() { + unregisterGlobal(API_NAME); + } + + public createLocal(id: string): LocalsKey { + return this.#getManager().createLocal(id); + } + + public getLocal(key: LocalsKey): T | undefined { + return this.#getManager().getLocal(key); + } + + public setLocal(key: LocalsKey, value: T): void { + return this.#getManager().setLocal(key, value); + } + + #getManager(): LocalsManager { + return getGlobal(API_NAME) ?? NOOP_LOCALS_MANAGER; + } +} diff --git a/packages/core/src/v3/locals/manager.ts b/packages/core/src/v3/locals/manager.ts new file mode 100644 index 0000000000..befe219260 --- /dev/null +++ b/packages/core/src/v3/locals/manager.ts @@ -0,0 +1,37 @@ +import { LocalsKey, LocalsManager } from "./types.js"; + +export class NoopLocalsManager implements LocalsManager { + createLocal(id: string): LocalsKey { + return { + __type: Symbol(), + id, + } as unknown as LocalsKey; + } + + getLocal(key: LocalsKey): T | undefined { + return undefined; + } + + setLocal(key: LocalsKey, value: T): void {} +} + +export class StandardLocalsManager implements LocalsManager { + private store: Map = new Map(); + + createLocal(id: string): LocalsKey { + const key = Symbol.for(id); + return { + __type: key, + id, + } as unknown as LocalsKey; + } + + getLocal(key: LocalsKey): T | undefined { + return this.store.get(key.__type) as T | undefined; + } + + setLocal(key: LocalsKey, value: T): void { + this.store.set(key.__type, value); + } +} +0; diff --git a/packages/core/src/v3/locals/types.ts b/packages/core/src/v3/locals/types.ts new file mode 100644 index 0000000000..aab683df09 --- /dev/null +++ b/packages/core/src/v3/locals/types.ts @@ -0,0 +1,14 @@ +declare const __local: unique symbol; +type BrandLocal = { [__local]: T }; + +// Create a type-safe store for your locals +export type LocalsKey = BrandLocal & { + readonly id: string; + readonly __type: unique symbol; +}; + +export interface LocalsManager { + createLocal(id: string): LocalsKey; + getLocal(key: LocalsKey): T | undefined; + setLocal(key: LocalsKey, value: T): void; +} diff --git a/packages/core/src/v3/runtime/managedRuntimeManager.ts b/packages/core/src/v3/runtime/managedRuntimeManager.ts index 67ef064498..b876a87084 100644 --- a/packages/core/src/v3/runtime/managedRuntimeManager.ts +++ b/packages/core/src/v3/runtime/managedRuntimeManager.ts @@ -1,3 +1,4 @@ +import { lifecycleHooks } from "../lifecycle-hooks-api.js"; import { BatchTaskRunExecutionResult, CompletedWaitpoint, @@ -44,9 +45,19 @@ export class ManagedRuntimeManager implements RuntimeManager { this.resolversByWaitId.set(params.id, resolve); }); + await lifecycleHooks.callOnWaitHookListeners({ + type: "task", + runId: params.id, + }); + const waitpoint = await promise; const result = this.waitpointToTaskRunExecutionResult(waitpoint); + await lifecycleHooks.callOnResumeHookListeners({ + type: "task", + runId: params.id, + }); + return result; }); } @@ -70,8 +81,20 @@ export class ManagedRuntimeManager implements RuntimeManager { }) ); + await lifecycleHooks.callOnWaitHookListeners({ + type: "batch", + batchId: params.id, + runCount: params.runCount, + }); + const waitpoints = await promise; + await lifecycleHooks.callOnResumeHookListeners({ + type: "batch", + batchId: params.id, + runCount: params.runCount, + }); + return { id: params.id, items: waitpoints.map(this.waitpointToTaskRunExecutionResult), @@ -91,8 +114,32 @@ export class ManagedRuntimeManager implements RuntimeManager { this.resolversByWaitId.set(waitpointFriendlyId, resolve); }); + if (finishDate) { + await lifecycleHooks.callOnWaitHookListeners({ + type: "duration", + date: finishDate, + }); + } else { + await lifecycleHooks.callOnWaitHookListeners({ + type: "token", + token: waitpointFriendlyId, + }); + } + const waitpoint = await promise; + if (finishDate) { + await lifecycleHooks.callOnResumeHookListeners({ + type: "duration", + date: finishDate, + }); + } else { + await lifecycleHooks.callOnResumeHookListeners({ + type: "token", + token: waitpointFriendlyId, + }); + } + return { ok: !waitpoint.outputIsError, output: waitpoint.output, diff --git a/packages/core/src/v3/schemas/build.ts b/packages/core/src/v3/schemas/build.ts index 0b122af2ed..c3df04eaa7 100644 --- a/packages/core/src/v3/schemas/build.ts +++ b/packages/core/src/v3/schemas/build.ts @@ -38,6 +38,7 @@ export const BuildManifest = z.object({ indexWorkerEntryPoint: z.string(), // Dev & Deploy has a indexWorkerEntryPoint indexControllerEntryPoint: z.string().optional(), // Only deploy has a indexControllerEntryPoint loaderEntryPoint: z.string().optional(), + initEntryPoint: z.string().optional(), // Optional init.ts entry point configPath: z.string(), externals: BuildExternal.array().optional(), build: z.object({ @@ -85,6 +86,7 @@ export const WorkerManifest = z.object({ workerEntryPoint: z.string(), controllerEntryPoint: z.string().optional(), loaderEntryPoint: z.string().optional(), + initEntryPoint: z.string().optional(), // Optional init.ts entry point runtime: BuildRuntime, customConditions: z.array(z.string()).optional(), otelImportHook: z diff --git a/packages/core/src/v3/schemas/common.ts b/packages/core/src/v3/schemas/common.ts index 43030847dc..5467603e9d 100644 --- a/packages/core/src/v3/schemas/common.ts +++ b/packages/core/src/v3/schemas/common.ts @@ -165,6 +165,7 @@ export const TaskRunInternalError = z.object({ "TASK_RUN_CANCELLED", "TASK_INPUT_ERROR", "TASK_OUTPUT_ERROR", + "TASK_MIDDLEWARE_ERROR", "HANDLE_ERROR_ERROR", "GRACEFUL_EXIT_TIMEOUT", "TASK_RUN_HEARTBEAT_TIMEOUT", diff --git a/packages/core/src/v3/semanticInternalAttributes.ts b/packages/core/src/v3/semanticInternalAttributes.ts index 63549bea25..4cfc03e8ec 100644 --- a/packages/core/src/v3/semanticInternalAttributes.ts +++ b/packages/core/src/v3/semanticInternalAttributes.ts @@ -33,6 +33,7 @@ export const SemanticInternalAttributes = { STYLE_ICON: "$style.icon", STYLE_VARIANT: "$style.variant", STYLE_ACCESSORY: "$style.accessory", + COLLAPSED: "$collapsed", METADATA: "$metadata", TRIGGER: "$trigger", PAYLOAD: "$payload", diff --git a/packages/core/src/v3/tracer.ts b/packages/core/src/v3/tracer.ts index 085107c017..42bc2f249e 100644 --- a/packages/core/src/v3/tracer.ts +++ b/packages/core/src/v3/tracer.ts @@ -168,7 +168,7 @@ export class TriggerTracer { const attributes = options?.attributes ?? {}; - const span = this.tracer.startSpan(name, options, ctx); + const span = this.tracer.startSpan(name, options, parentContext); this.tracer .startSpan( diff --git a/packages/core/src/v3/types/tasks.ts b/packages/core/src/v3/types/tasks.ts index e8e54771c4..8c1e4d1014 100644 --- a/packages/core/src/v3/types/tasks.ts +++ b/packages/core/src/v3/types/tasks.ts @@ -1,5 +1,18 @@ import { SerializableJson } from "../../schemas/json.js"; import { TriggerApiRequestOptions } from "../apiClient/index.js"; +import { + AnyOnCatchErrorHookFunction, + OnCatchErrorHookFunction, + OnCleanupHookFunction, + OnCompleteHookFunction, + OnFailureHookFunction, + OnInitHookFunction, + OnMiddlewareHookFunction, + OnResumeHookFunction, + OnStartHookFunction, + OnSuccessHookFunction, + OnWaitHookFunction, +} from "../lifecycleHooks/types.js"; import { RunTags } from "../schemas/api.js"; import { MachineCpu, @@ -10,10 +23,10 @@ import { TaskRunContext, } from "../schemas/index.js"; import { IdempotencyKey } from "./idempotencyKeys.js"; +import { QueueOptions } from "./queues.js"; import { AnySchemaParseFn, inferSchemaIn, inferSchemaOut, Schema } from "./schemas.js"; -import { Prettify } from "./utils.js"; import { inferToolParameters, ToolTaskParameters } from "./tools.js"; -import { QueueOptions } from "./queues.js"; +import { Prettify } from "./utils.js"; export type Queue = QueueOptions; export type TaskSchema = Schema; @@ -94,6 +107,7 @@ export type InitFnParams = Prettify<{ export type StartFnParams = Prettify<{ ctx: Context; + init?: InitOutput; /** Abort signal that is aborted when a task run exceeds it's maxDuration. Can be used to automatically cancel downstream requests */ signal?: AbortSignal; }>; @@ -127,7 +141,7 @@ export type HandleErrorResult = export type HandleErrorArgs = { ctx: Context; - init: unknown; + init?: Record; retry?: RetryOptions; retryAt?: Date; retryDelayInMs?: number; @@ -258,22 +272,33 @@ type CommonTaskOptions< /** * init is called before the run function is called. It's useful for setting up any global state. + * + * @deprecated Use locals and middleware instead */ - init?: (payload: TPayload, params: InitFnParams) => Promise; + init?: OnInitHookFunction; /** * cleanup is called after the run function has completed. + * + * @deprecated Use middleware instead */ - cleanup?: (payload: TPayload, params: RunFnParams) => Promise; + cleanup?: OnCleanupHookFunction; /** * handleError is called when the run function throws an error. It can be used to modify the error or return new retry options. + * + * @deprecated Use catchError instead + */ + handleError?: OnCatchErrorHookFunction; + + /** + * catchError is called when the run function throws an error. It can be used to modify the error or return new retry options. */ - handleError?: ( - payload: TPayload, - error: unknown, - params: HandleErrorFnParams - ) => HandleErrorResult; + catchError?: OnCatchErrorHookFunction; + + onResume?: OnResumeHookFunction; + onWait?: OnWaitHookFunction; + onComplete?: OnCompleteHookFunction; /** * middleware allows you to run code "around" the run function. This can be useful for logging, metrics, or other cross-cutting concerns. @@ -292,30 +317,22 @@ type CommonTaskOptions< * }); * ``` */ - middleware?: (payload: TPayload, params: MiddlewareFnParams) => Promise; + middleware?: OnMiddlewareHookFunction; /** * onStart is called the first time a task is executed in a run (not before every retry) */ - onStart?: (payload: TPayload, params: StartFnParams) => Promise; + onStart?: OnStartHookFunction; /** * onSuccess is called after the run function has successfully completed. */ - onSuccess?: ( - payload: TPayload, - output: TOutput, - params: SuccessFnParams - ) => Promise; + onSuccess?: OnSuccessHookFunction; /** * onFailure is called after a task run has failed (meaning the run function threw an error and won't be retried anymore) */ - onFailure?: ( - payload: TPayload, - error: unknown, - params: FailureFnParams - ) => Promise; + onFailure?: OnFailureHookFunction; }; export type TaskOptions< diff --git a/packages/core/src/v3/utils/globals.ts b/packages/core/src/v3/utils/globals.ts index f3ec155751..e59539b343 100644 --- a/packages/core/src/v3/utils/globals.ts +++ b/packages/core/src/v3/utils/globals.ts @@ -1,5 +1,7 @@ import { ApiClientConfiguration } from "../apiClientManager/types.js"; import { Clock } from "../clock/clock.js"; +import { LifecycleHooksManager } from "../lifecycleHooks/types.js"; +import { LocalsManager } from "../locals/types.js"; import { ResourceCatalog } from "../resource-catalog/catalog.js"; import { RunMetadataManager } from "../runMetadata/types.js"; import type { RuntimeManager } from "../runtime/manager.js"; @@ -62,4 +64,6 @@ type TriggerDotDevGlobalAPI = { ["timeout"]?: TimeoutManager; ["wait-until"]?: WaitUntilManager; ["run-timeline-metrics"]?: RunTimelineMetricsManager; + ["lifecycle-hooks"]?: LifecycleHooksManager; + ["locals"]?: LocalsManager; }; diff --git a/packages/core/src/v3/workers/index.ts b/packages/core/src/v3/workers/index.ts index 76cdb7110f..7d67a23836 100644 --- a/packages/core/src/v3/workers/index.ts +++ b/packages/core/src/v3/workers/index.ts @@ -20,3 +20,5 @@ export { ManagedRuntimeManager } from "../runtime/managedRuntimeManager.js"; export * from "../runEngineWorker/index.js"; export { StandardRunTimelineMetricsManager } from "../runTimelineMetrics/runTimelineMetricsManager.js"; export { WarmStartClient, type WarmStartClientOptions } from "../workers/warmStartClient.js"; +export { StandardLifecycleHooksManager } from "../lifecycleHooks/manager.js"; +export { StandardLocalsManager } from "../locals/manager.js"; diff --git a/packages/core/src/v3/workers/taskExecutor.ts b/packages/core/src/v3/workers/taskExecutor.ts index a8339903ba..355dbec95c 100644 --- a/packages/core/src/v3/workers/taskExecutor.ts +++ b/packages/core/src/v3/workers/taskExecutor.ts @@ -2,11 +2,25 @@ import { SpanKind } from "@opentelemetry/api"; import { VERSION } from "../../version.js"; import { ApiError, RateLimitError } from "../apiClient/errors.js"; import { ConsoleInterceptor } from "../consoleInterceptor.js"; -import { isInternalError, parseError, sanitizeError, TaskPayloadParsedError } from "../errors.js"; -import { runMetadata, TriggerConfig, waitUntil } from "../index.js"; +import { + InternalError, + isInternalError, + parseError, + sanitizeError, + TaskPayloadParsedError, +} from "../errors.js"; +import { flattenAttributes, lifecycleHooks, runMetadata, waitUntil } from "../index.js"; +import { + AnyOnMiddlewareHookFunction, + RegisteredHookFunction, + TaskCompleteResult, + TaskInitOutput, + TaskWait, +} from "../lifecycleHooks/types.js"; import { recordSpanException, TracingSDK } from "../otel/index.js"; import { runTimelineMetrics } from "../run-timeline-metrics-api.js"; import { + RetryOptions, ServerBackgroundWorker, TaskRunContext, TaskRunErrorCodes, @@ -17,8 +31,8 @@ import { import { SemanticInternalAttributes } from "../semanticInternalAttributes.js"; import { taskContext } from "../task-context-api.js"; import { TriggerTracer } from "../tracer.js"; -import { HandleErrorFunction, TaskMetadataWithFunctions } from "../types/index.js"; -import { UsageMeasurement } from "../usage/types.js"; +import { tryCatch } from "../tryCatch.js"; +import { HandleErrorModificationOptions, TaskMetadataWithFunctions } from "../types/index.js"; import { conditionallyExportPacket, conditionallyImportPacket, @@ -32,16 +46,22 @@ export type TaskExecutorOptions = { tracingSDK: TracingSDK; tracer: TriggerTracer; consoleInterceptor: ConsoleInterceptor; - config: TriggerConfig | undefined; - handleErrorFn: HandleErrorFunction | undefined; + retries?: { + enabledInDev?: boolean; + default?: RetryOptions; + }; }; export class TaskExecutor { private _tracingSDK: TracingSDK; private _tracer: TriggerTracer; private _consoleInterceptor: ConsoleInterceptor; - private _importedConfig: TriggerConfig | undefined; - private _handleErrorFn: HandleErrorFunction | undefined; + private _retries: + | { + enabledInDev?: boolean; + default?: RetryOptions; + } + | undefined; constructor( public task: TaskMetadataWithFunctions, @@ -50,15 +70,13 @@ export class TaskExecutor { this._tracingSDK = options.tracingSDK; this._tracer = options.tracer; this._consoleInterceptor = options.consoleInterceptor; - this._importedConfig = options.config; - this._handleErrorFn = options.handleErrorFn; + this._retries = options.retries; } async execute( execution: TaskRunExecution, worker: ServerBackgroundWorker, traceContext: Record, - usage: UsageMeasurement, signal?: AbortSignal ): Promise<{ result: TaskRunExecutionResult }> { const ctx = TaskRunContext.parse(execution); @@ -91,110 +109,81 @@ export class TaskExecutor { let parsedPayload: any; let initOutput: any; - try { - await runTimelineMetrics.measureMetric("trigger.dev/execution", "payload", async () => { + const [inputError, payloadResult] = await tryCatch( + runTimelineMetrics.measureMetric("trigger.dev/execution", "payload", async () => { const payloadPacket = await conditionallyImportPacket(originalPacket, this._tracer); - parsedPayload = await parsePacket(payloadPacket); - }); - } catch (inputError) { - recordSpanException(span, inputError); + return await parsePacket(payloadPacket); + }) + ); - return { - ok: false, - id: execution.run.id, - error: { - type: "INTERNAL_ERROR", - code: TaskRunErrorCodes.TASK_INPUT_ERROR, - message: - inputError instanceof Error - ? `${inputError.name}: ${inputError.message}` - : typeof inputError === "string" - ? inputError - : undefined, - stackTrace: inputError instanceof Error ? inputError.stack : undefined, - }, - } satisfies TaskRunExecutionResult; + if (inputError) { + recordSpanException(span, inputError); + return this.#internalErrorResult( + execution, + TaskRunErrorCodes.TASK_INPUT_ERROR, + inputError + ); } - try { - parsedPayload = await this.#parsePayload(parsedPayload); - - if (execution.attempt.number === 1) { - await this.#callOnStartFunctions(parsedPayload, ctx, signal); - } + parsedPayload = await this.#parsePayload(payloadResult); - initOutput = await this.#callInitFunctions(parsedPayload, ctx, signal); + lifecycleHooks.registerOnWaitHookListener(async (wait) => { + await this.#callOnWaitFunctions(wait, parsedPayload, ctx, initOutput, signal); + }); - const output = await this.#callRun(parsedPayload, ctx, initOutput, signal); + lifecycleHooks.registerOnResumeHookListener(async (wait) => { + await this.#callOnResumeFunctions(wait, parsedPayload, ctx, initOutput, signal); + }); - await this.#callOnSuccessFunctions(parsedPayload, output, ctx, initOutput, signal); + const executeTask = async (payload: any) => { + const [runError, output] = await tryCatch( + (async () => { + initOutput = await this.#callInitFunctions(payload, ctx, signal); - try { - const stringifiedOutput = await stringifyIO(output); + if (execution.attempt.number === 1) { + await this.#callOnStartFunctions(payload, ctx, initOutput, signal); + } - const finalOutput = await conditionallyExportPacket( - stringifiedOutput, - `${execution.attempt.id}/output`, - this._tracer - ); + return await this.#callRun(payload, ctx, initOutput, signal); + })() + ); - const attributes = await createPacketAttributes( - finalOutput, - SemanticInternalAttributes.OUTPUT, - SemanticInternalAttributes.OUTPUT_TYPE + if (runError) { + const [handleErrorError, handleErrorResult] = await tryCatch( + this.#handleError(execution, runError, payload, ctx, initOutput, signal) ); - if (attributes) { - span.setAttributes(attributes); + if (handleErrorError) { + recordSpanException(span, handleErrorError); + return this.#internalErrorResult( + execution, + TaskRunErrorCodes.HANDLE_ERROR_ERROR, + handleErrorError + ); } - return { - ok: true, - id: execution.run.id, - output: finalOutput.data, - outputType: finalOutput.dataType, - } satisfies TaskRunExecutionResult; - } catch (outputError) { - recordSpanException(span, outputError); - - return { - ok: false, - id: execution.run.id, - error: { - type: "INTERNAL_ERROR", - code: TaskRunErrorCodes.TASK_OUTPUT_ERROR, - message: - outputError instanceof Error - ? outputError.message - : typeof outputError === "string" - ? outputError - : undefined, - }, - } satisfies TaskRunExecutionResult; - } - } catch (runError) { - try { - const handleErrorResult = await this.#handleError( - execution, - runError, - parsedPayload, - ctx, - initOutput, - signal - ); - recordSpanException(span, handleErrorResult.error ?? runError); if (handleErrorResult.status !== "retry") { await this.#callOnFailureFunctions( - parsedPayload, + payload, handleErrorResult.error ?? runError, ctx, initOutput, signal ); + + await this.#callOnCompleteFunctions( + payload, + { ok: false, error: handleErrorResult.error ?? runError }, + ctx, + initOutput, + signal + ); } + await this.#cleanupAndWaitUntil(payload, ctx, initOutput, signal); + return { id: execution.run.id, ok: false, @@ -206,30 +195,87 @@ export class TaskExecutor { retry: handleErrorResult.status === "retry" ? handleErrorResult.retry : undefined, skippedRetrying: handleErrorResult.status === "skipped", } satisfies TaskRunExecutionResult; - } catch (handleErrorError) { - recordSpanException(span, handleErrorError); + } - return { - ok: false, - id: execution.run.id, - error: { - type: "INTERNAL_ERROR", - code: TaskRunErrorCodes.HANDLE_ERROR_ERROR, - message: - handleErrorError instanceof Error - ? handleErrorError.message - : typeof handleErrorError === "string" - ? handleErrorError - : undefined, - }, - } satisfies TaskRunExecutionResult; + const [outputError, stringifiedOutput] = await tryCatch(stringifyIO(output)); + + if (outputError) { + recordSpanException(span, outputError); + await this.#cleanupAndWaitUntil(payload, ctx, initOutput, signal); + + return this.#internalErrorResult( + execution, + TaskRunErrorCodes.TASK_OUTPUT_ERROR, + outputError + ); } - } finally { - await this.#callTaskCleanup(parsedPayload, ctx, initOutput, signal); - await this.#blockForWaitUntil(); - span.setAttributes(runTimelineMetrics.convertMetricsToSpanAttributes()); - } + const [exportError, finalOutput] = await tryCatch( + conditionallyExportPacket( + stringifiedOutput, + `${execution.attempt.id}/output`, + this._tracer + ) + ); + + if (exportError) { + recordSpanException(span, exportError); + await this.#cleanupAndWaitUntil(payload, ctx, initOutput, signal); + + return this.#internalErrorResult( + execution, + TaskRunErrorCodes.TASK_OUTPUT_ERROR, + exportError + ); + } + + const [attrError, attributes] = await tryCatch( + createPacketAttributes( + finalOutput, + SemanticInternalAttributes.OUTPUT, + SemanticInternalAttributes.OUTPUT_TYPE + ) + ); + + if (!attrError && attributes) { + span.setAttributes(attributes); + } + + await this.#callOnSuccessFunctions(payload, output, ctx, initOutput, signal); + await this.#callOnCompleteFunctions( + payload, + { ok: true, data: output }, + ctx, + initOutput, + signal + ); + + await this.#cleanupAndWaitUntil(payload, ctx, initOutput, signal); + + return { + ok: true, + id: execution.run.id, + output: finalOutput.data, + outputType: finalOutput.dataType, + } satisfies TaskRunExecutionResult; + }; + + const globalMiddlewareHooks = lifecycleHooks.getGlobalMiddlewareHooks(); + const taskMiddlewareHook = lifecycleHooks.getTaskMiddlewareHook(this.task.id); + + const middlewareHooks = [ + ...globalMiddlewareHooks, + taskMiddlewareHook ? { id: this.task.id, fn: taskMiddlewareHook } : undefined, + ].filter(Boolean) as RegisteredHookFunction[]; + + return await this.#executeTaskWithMiddlewareHooks( + parsedPayload, + ctx, + execution, + middlewareHooks, + executeTask, + signal + ); }); }, { @@ -253,196 +299,479 @@ export class TaskExecutor { return { result }; } + async #executeTaskWithMiddlewareHooks( + payload: unknown, + ctx: TaskRunContext, + execution: TaskRunExecution, + hooks: RegisteredHookFunction[], + executeTask: (payload: unknown) => Promise, + signal?: AbortSignal + ) { + let output: any; + let executeError: unknown; + + const runner = hooks.reduceRight( + (next, hook) => { + return async () => { + await this._tracer.startActiveSpan( + hook.name ? `middleware/${hook.name}` : "middleware", + async (span) => { + await hook.fn({ payload, ctx, signal, task: this.task.id, next }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-middleware", + }, + } + ); + }; + }, + async () => { + const [error, result] = await tryCatch(executeTask(payload)); + if (error) { + executeError = error; + } else { + output = result; + } + } + ); + + const [runnerError] = await tryCatch(runner()); + if (runnerError) { + return this.#internalErrorResult( + execution, + TaskRunErrorCodes.TASK_MIDDLEWARE_ERROR, + runnerError + ); + } + + if (executeError) { + throw executeError; + } + + return output; + } + async #callRun(payload: unknown, ctx: TaskRunContext, init: unknown, signal?: AbortSignal) { const runFn = this.task.fns.run; - const middlewareFn = this.task.fns.middleware; if (!runFn) { throw new Error("Task does not have a run function"); } - if (!middlewareFn) { - return runTimelineMetrics.measureMetric("trigger.dev/execution", "run", () => - runFn(payload, { ctx, init, signal }) - ); - } + // Create a promise that rejects when the signal aborts + const abortPromise = signal + ? new Promise((_, reject) => { + signal.addEventListener("abort", () => { + const maxDuration = ctx.run.maxDuration; + reject( + new InternalError({ + code: TaskRunErrorCodes.MAX_DURATION_EXCEEDED, + message: `Task execution exceeded maximum duration of ${maxDuration}ms`, + }) + ); + }); + }) + : undefined; + + return runTimelineMetrics.measureMetric("trigger.dev/execution", "run", async () => { + return await this._tracer.startActiveSpan( + "run", + async (span) => { + if (abortPromise) { + // Race between the run function and the abort promise + return await Promise.race([runFn(payload, { ctx, init, signal }), abortPromise]); + } - return middlewareFn(payload, { - ctx, - signal, - next: async () => - runTimelineMetrics.measureMetric("trigger.dev/execution", "run", () => - runFn(payload, { ctx, init, signal }) - ), + return await runFn(payload, { ctx, init, signal }); + }, + { + attributes: { [SemanticInternalAttributes.STYLE_ICON]: "task-fn-run" }, + } + ); }); } - async #callInitFunctions(payload: unknown, ctx: TaskRunContext, signal?: AbortSignal) { - await this.#callConfigInit(payload, ctx, signal); - - const initFn = this.task.fns.init; + async #callOnWaitFunctions( + wait: TaskWait, + payload: unknown, + ctx: TaskRunContext, + initOutput: TaskInitOutput, + signal?: AbortSignal + ) { + const globalWaitHooks = lifecycleHooks.getGlobalWaitHooks(); + const taskWaitHook = lifecycleHooks.getTaskWaitHook(this.task.id); - if (!initFn) { - return {}; + if (globalWaitHooks.length === 0 && !taskWaitHook) { + return; } - return this._tracer.startActiveSpan( - "init", - async (span) => { - return await runTimelineMetrics.measureMetric("trigger.dev/execution", "init", () => - initFn(payload, { ctx, signal }) - ); - }, - { - attributes: { - [SemanticInternalAttributes.STYLE_ICON]: "function", - }, - } - ); - } + const result = await runTimelineMetrics.measureMetric( + "trigger.dev/execution", + "onWait", + async () => { + for (const hook of globalWaitHooks) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + hook.name ? `onWait/${hook.name}` : "onWait/global", + async (span) => { + await hook.fn({ payload, ctx, signal, task: this.task.id, wait, init: initOutput }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onWait", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); - async #callConfigInit(payload: unknown, ctx: TaskRunContext, signal?: AbortSignal) { - const initFn = this._importedConfig?.init; + if (hookError) { + throw hookError; + } + } - if (!initFn) { - return {}; - } + if (taskWaitHook) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + "onWait/task", + async (span) => { + await taskWaitHook({ + payload, + ctx, + signal, + task: this.task.id, + wait, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onWait", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); - return this._tracer.startActiveSpan( - "config.init", - async (span) => { - return await runTimelineMetrics.measureMetric( - "trigger.dev/execution", - "config.init", - async () => initFn(payload, { ctx, signal }) - ); - }, - { - attributes: { - [SemanticInternalAttributes.STYLE_ICON]: "function", - }, + if (hookError) { + throw hookError; + } + } } ); } - async #callOnSuccessFunctions( + async #callOnResumeFunctions( + wait: TaskWait, payload: unknown, - output: any, ctx: TaskRunContext, - initOutput: any, + initOutput: TaskInitOutput, signal?: AbortSignal ) { - await this.#callOnSuccessFunction( - this.task.fns.onSuccess, - "task.onSuccess", - payload, - output, - ctx, - initOutput, - signal - ); + const globalResumeHooks = lifecycleHooks.getGlobalResumeHooks(); + const taskResumeHook = lifecycleHooks.getTaskResumeHook(this.task.id); - await this.#callOnSuccessFunction( - this._importedConfig?.onSuccess, - "config.onSuccess", - payload, - output, - ctx, - initOutput, - signal + if (globalResumeHooks.length === 0 && !taskResumeHook) { + return; + } + + const result = await runTimelineMetrics.measureMetric( + "trigger.dev/execution", + "onResume", + async () => { + for (const hook of globalResumeHooks) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + hook.name ? `onResume/${hook.name}` : "onResume/global", + async (span) => { + await hook.fn({ payload, ctx, signal, task: this.task.id, wait, init: initOutput }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onResume", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; + } + } + + if (taskResumeHook) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + "onResume/task", + async (span) => { + await taskResumeHook({ + payload, + ctx, + signal, + task: this.task.id, + wait, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onResume", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; + } + } + } ); } - async #callOnSuccessFunction( - onSuccessFn: TaskMetadataWithFunctions["fns"]["onSuccess"], - name: string, - payload: unknown, - output: any, - ctx: TaskRunContext, - initOutput: any, - signal?: AbortSignal - ) { - if (!onSuccessFn) { - return; + async #callInitFunctions(payload: unknown, ctx: TaskRunContext, signal?: AbortSignal) { + const globalInitHooks = lifecycleHooks.getGlobalInitHooks(); + const taskInitHook = lifecycleHooks.getTaskInitHook(this.task.id); + + if (globalInitHooks.length === 0 && !taskInitHook) { + return {}; } - try { - await this._tracer.startActiveSpan( - name, - async (span) => { - return await runTimelineMetrics.measureMetric("trigger.dev/execution", name, () => - onSuccessFn(payload, output, { ctx, init: initOutput, signal }) + const result = await runTimelineMetrics.measureMetric( + "trigger.dev/execution", + "init", + async () => { + // Store global hook results in an array + const globalResults = []; + for (const hook of globalInitHooks) { + const [hookError, result] = await tryCatch( + this._tracer.startActiveSpan( + hook.name ? `init/${hook.name}` : "init/global", + async (span) => { + const result = await hook.fn({ payload, ctx, signal, task: this.task.id }); + + if (result && typeof result === "object" && !Array.isArray(result)) { + span.setAttributes(flattenAttributes(result)); + return result; + } + + return {}; + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-init", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) ); - }, - { - attributes: { - [SemanticInternalAttributes.STYLE_ICON]: "function", - }, + + if (hookError) { + throw hookError; + } + + if (result && typeof result === "object" && !Array.isArray(result)) { + globalResults.push(result); + } } - ); - } catch { - // Ignore errors from onSuccess functions + + // Merge all global results into a single object + const mergedGlobalResults = Object.assign({}, ...globalResults); + + if (taskInitHook) { + const [hookError, taskResult] = await tryCatch( + this._tracer.startActiveSpan( + "init/task", + async (span) => { + const result = await taskInitHook({ payload, ctx, signal, task: this.task.id }); + + if (result && typeof result === "object" && !Array.isArray(result)) { + span.setAttributes(flattenAttributes(result)); + return result; + } + + return {}; + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-init", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; + } + + // Only merge if taskResult is an object + if (taskResult && typeof taskResult === "object" && !Array.isArray(taskResult)) { + return { ...mergedGlobalResults, ...taskResult }; + } + + // If taskResult isn't an object, return global results + return mergedGlobalResults; + } + + return mergedGlobalResults; + } + ); + + if (result && typeof result === "object" && !Array.isArray(result)) { + return result; } + + return; } - async #callOnFailureFunctions( + async #callOnSuccessFunctions( payload: unknown, - error: unknown, + output: any, ctx: TaskRunContext, initOutput: any, signal?: AbortSignal ) { - await this.#callOnFailureFunction( - this.task.fns.onFailure, - "task.onFailure", - payload, - error, - ctx, - initOutput, - signal - ); + const globalSuccessHooks = lifecycleHooks.getGlobalSuccessHooks(); + const taskSuccessHook = lifecycleHooks.getTaskSuccessHook(this.task.id); - await this.#callOnFailureFunction( - this._importedConfig?.onFailure, - "config.onFailure", - payload, - error, - ctx, - initOutput, - signal - ); + if (globalSuccessHooks.length === 0 && !taskSuccessHook) { + return; + } + + return await runTimelineMetrics.measureMetric("trigger.dev/execution", "success", async () => { + for (const hook of globalSuccessHooks) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + hook.name ? `onSuccess/${hook.name}` : "onSuccess/global", + async (span) => { + await hook.fn({ + payload, + output, + ctx, + signal, + task: this.task.id, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onSuccess", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; + } + } + + if (taskSuccessHook) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + "onSuccess/task", + async (span) => { + await taskSuccessHook({ + payload, + output, + ctx, + signal, + task: this.task.id, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onSuccess", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; + } + } + }); } - async #callOnFailureFunction( - onFailureFn: TaskMetadataWithFunctions["fns"]["onFailure"], - name: string, + async #callOnFailureFunctions( payload: unknown, error: unknown, ctx: TaskRunContext, initOutput: any, signal?: AbortSignal ) { - if (!onFailureFn) { + const globalFailureHooks = lifecycleHooks.getGlobalFailureHooks(); + const taskFailureHook = lifecycleHooks.getTaskFailureHook(this.task.id); + + if (globalFailureHooks.length === 0 && !taskFailureHook) { return; } - try { - return await this._tracer.startActiveSpan( - name, - async (span) => { - return await runTimelineMetrics.measureMetric("trigger.dev/execution", name, () => - onFailureFn(payload, error, { ctx, init: initOutput, signal }) - ); - }, - { - attributes: { - [SemanticInternalAttributes.STYLE_ICON]: "function", - }, + return await runTimelineMetrics.measureMetric("trigger.dev/execution", "failure", async () => { + for (const hook of globalFailureHooks) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + hook.name ? `onFailure/${hook.name}` : "onFailure/global", + async (span) => { + await hook.fn({ + payload, + error, + ctx, + signal, + task: this.task.id, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onFailure", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; } - ); - } catch (e) { - // Ignore errors from onFailure functions - } + } + + if (taskFailureHook) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + "onFailure/task", + async (span) => { + await taskFailureHook({ + payload, + error, + ctx, + signal, + task: this.task.id, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onFailure", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; + } + } + }); } async #parsePayload(payload: unknown) { @@ -450,78 +779,154 @@ export class TaskExecutor { return payload; } - try { - return await this.task.fns.parsePayload(payload); - } catch (e) { - throw new TaskPayloadParsedError(e); + const [parseError, result] = await tryCatch(this.task.fns.parsePayload(payload)); + if (parseError) { + throw new TaskPayloadParsedError(parseError); } + return result; } - async #callOnStartFunctions(payload: unknown, ctx: TaskRunContext, signal?: AbortSignal) { - await this.#callOnStartFunction( - this._importedConfig?.onStart, - "config.onStart", - payload, - ctx, - {}, - signal - ); - - await this.#callOnStartFunction( - this.task.fns.onStart, - "task.onStart", - payload, - ctx, - {}, - signal - ); - } - - async #callOnStartFunction( - onStartFn: TaskMetadataWithFunctions["fns"]["onStart"], - name: string, + async #callOnStartFunctions( payload: unknown, ctx: TaskRunContext, initOutput: any, signal?: AbortSignal ) { - if (!onStartFn) { + const globalStartHooks = lifecycleHooks.getGlobalStartHooks(); + const taskStartHook = lifecycleHooks.getTaskStartHook(this.task.id); + + if (globalStartHooks.length === 0 && !taskStartHook) { return; } - try { - await this._tracer.startActiveSpan( - name, - async (span) => { - return await runTimelineMetrics.measureMetric("trigger.dev/execution", name, () => - onStartFn(payload, { ctx, signal }) - ); - }, - { - attributes: { - [SemanticInternalAttributes.STYLE_ICON]: "function", - }, + return await runTimelineMetrics.measureMetric("trigger.dev/execution", "start", async () => { + for (const hook of globalStartHooks) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + hook.name ? `onStart/${hook.name}` : "onStart/global", + async (span) => { + await hook.fn({ payload, ctx, signal, task: this.task.id, init: initOutput }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; } - ); - } catch { - // Ignore errors from onStart functions - } + } + + if (taskStartHook) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + "onStart/task", + async (span) => { + await taskStartHook({ + payload, + ctx, + signal, + task: this.task.id, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; + } + } + }); } - async #callTaskCleanup( + async #cleanupAndWaitUntil( payload: unknown, ctx: TaskRunContext, - init: unknown, + initOutput: any, signal?: AbortSignal ) { - const cleanupFn = this.task.fns.cleanup; + await this.#callCleanupFunctions(payload, ctx, initOutput, signal); + await this.#blockForWaitUntil(); + } - if (!cleanupFn) { + async #callCleanupFunctions( + payload: unknown, + ctx: TaskRunContext, + initOutput: any, + signal?: AbortSignal + ) { + const globalCleanupHooks = lifecycleHooks.getGlobalCleanupHooks(); + const taskCleanupHook = lifecycleHooks.getTaskCleanupHook(this.task.id); + + if (globalCleanupHooks.length === 0 && !taskCleanupHook) { return; } - return this._tracer.startActiveSpan("cleanup", async (span) => { - return await cleanupFn(payload, { ctx, init, signal }); + return await runTimelineMetrics.measureMetric("trigger.dev/execution", "cleanup", async () => { + for (const hook of globalCleanupHooks) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + hook.name ? `cleanup/${hook.name}` : "cleanup/global", + async (span) => { + await hook.fn({ + payload, + ctx, + signal, + task: this.task.id, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-cleanup", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; + } + } + + if (taskCleanupHook) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + "cleanup/task", + async (span) => { + await taskCleanupHook({ + payload, + ctx, + signal, + task: this.task.id, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-cleanup", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; + } + } }); } @@ -538,6 +943,7 @@ export class TaskExecutor { { attributes: { [SemanticInternalAttributes.STYLE_ICON]: "clock", + [SemanticInternalAttributes.COLLAPSED]: true, }, } ); @@ -548,21 +954,17 @@ export class TaskExecutor { error: unknown, payload: any, ctx: TaskRunContext, - init: unknown, + init: TaskInitOutput, signal?: AbortSignal ): Promise< | { status: "retry"; retry: TaskRunExecutionRetry; error?: unknown } - | { status: "skipped"; error?: unknown } // skipped is different than noop, it means that the task was skipped from retrying, instead of just not retrying + | { status: "skipped"; error?: unknown } | { status: "noop"; error?: unknown } > { - const retriesConfig = this._importedConfig?.retries; - + const retriesConfig = this._retries; const retry = this.task.retry ?? retriesConfig?.default; - if (!retry) { - return { status: "noop" }; - } - + // Early exit conditions that prevent retrying if (isInternalError(error) && error.skipRetrying) { return { status: "skipped", error }; } @@ -574,23 +976,28 @@ export class TaskExecutor { return { status: "skipped" }; } - if (execution.run.maxAttempts) { - retry.maxAttempts = Math.max(execution.run.maxAttempts, 1); - } - - let delay = calculateNextRetryDelay(retry, execution.attempt.number); - - if ( - delay && - error instanceof Error && - error.name === "TriggerApiError" && - (error as ApiError).status === 429 - ) { - const rateLimitError = error as RateLimitError; + // Calculate default retry delay if retry config exists + let defaultDelay: number | undefined; + if (retry) { + if (execution.run.maxAttempts) { + retry.maxAttempts = Math.max(execution.run.maxAttempts, 1); + } - delay = rateLimitError.millisecondsUntilReset; + defaultDelay = calculateNextRetryDelay(retry, execution.attempt.number); + + // Handle rate limit errors + if ( + defaultDelay && + error instanceof Error && + error.name === "TriggerApiError" && + (error as ApiError).status === 429 + ) { + const rateLimitError = error as RateLimitError; + defaultDelay = rateLimitError.millisecondsUntilReset; + } } + // Check if retries are enabled in dev environment if ( execution.environment.type === "DEVELOPMENT" && typeof retriesConfig?.enabledInDev === "boolean" && @@ -600,80 +1007,203 @@ export class TaskExecutor { } return this._tracer.startActiveSpan( - "handleError()", + "catchError", async (span) => { - const handleErrorResult = this.task.fns.handleError - ? await this.task.fns.handleError(payload, error, { - ctx, - init, - retry, - retryDelayInMs: delay, - retryAt: delay ? new Date(Date.now() + delay) : undefined, - signal, - }) - : this._importedConfig - ? await this._handleErrorFn?.(payload, error, { - ctx, - init, - retry, - retryDelayInMs: delay, - retryAt: delay ? new Date(Date.now() + delay) : undefined, - signal, - }) - : undefined; - - // If handleErrorResult - if (!handleErrorResult) { - return typeof delay === "undefined" - ? { status: "noop" } - : { status: "retry", retry: { timestamp: Date.now() + delay, delay } }; + // Try task-specific catch error hook first + const taskCatchErrorHook = lifecycleHooks.getTaskCatchErrorHook(this.task.id); + if (taskCatchErrorHook) { + const result = await taskCatchErrorHook({ + payload, + error, + ctx, + init, + retry, + retryDelayInMs: defaultDelay, + retryAt: defaultDelay ? new Date(Date.now() + defaultDelay) : undefined, + signal, + task: this.task.id, + }); + + if (result) { + return this.#processHandleErrorResult(result, execution.attempt.number, defaultDelay); + } } - if (handleErrorResult.skipRetrying) { - return { status: "skipped", error: handleErrorResult.error }; + // Try global catch error hooks in order + const globalCatchErrorHooks = lifecycleHooks.getGlobalCatchErrorHooks(); + for (const hook of globalCatchErrorHooks) { + const result = await hook.fn({ + payload, + error, + ctx, + init, + retry, + retryDelayInMs: defaultDelay, + retryAt: defaultDelay ? new Date(Date.now() + defaultDelay) : undefined, + signal, + task: this.task.id, + }); + + if (result) { + return this.#processHandleErrorResult(result, execution.attempt.number, defaultDelay); + } } - if (typeof handleErrorResult.retryAt !== "undefined") { - return { + // If no hooks handled the error, use default retry behavior + return typeof defaultDelay === "undefined" + ? { status: "noop" as const } + : { + status: "retry" as const, + retry: { timestamp: Date.now() + defaultDelay, delay: defaultDelay }, + }; + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-catchError", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ); + } + + // Helper method to process handle error results + #processHandleErrorResult( + result: HandleErrorModificationOptions, + attemptNumber: number, + defaultDelay?: number + ): + | { status: "retry"; retry: TaskRunExecutionRetry; error?: unknown } + | { status: "skipped"; error?: unknown } + | { status: "noop"; error?: unknown } { + if (result.skipRetrying) { + return { status: "skipped", error: result.error }; + } + + if (typeof result.retryAt !== "undefined") { + return { + status: "retry", + retry: { + timestamp: result.retryAt.getTime(), + delay: result.retryAt.getTime() - Date.now(), + }, + error: result.error, + }; + } + + if (typeof result.retryDelayInMs === "number") { + return { + status: "retry", + retry: { + timestamp: Date.now() + result.retryDelayInMs, + delay: result.retryDelayInMs, + }, + error: result.error, + }; + } + + if (result.retry && typeof result.retry === "object") { + const delay = calculateNextRetryDelay(result.retry, attemptNumber); + + return typeof delay === "undefined" + ? { status: "noop", error: result.error } + : { status: "retry", - retry: { - timestamp: handleErrorResult.retryAt.getTime(), - delay: handleErrorResult.retryAt.getTime() - Date.now(), - }, - error: handleErrorResult.error, + retry: { timestamp: Date.now() + delay, delay }, + error: result.error, }; - } + } - if (typeof handleErrorResult.retryDelayInMs === "number") { - return { - status: "retry", - retry: { - timestamp: Date.now() + handleErrorResult.retryDelayInMs, - delay: handleErrorResult.retryDelayInMs, + return { status: "noop", error: result.error }; + } + + async #callOnCompleteFunctions( + payload: unknown, + result: TaskCompleteResult, + ctx: TaskRunContext, + initOutput: any, + signal?: AbortSignal + ) { + const globalCompleteHooks = lifecycleHooks.getGlobalCompleteHooks(); + const taskCompleteHook = lifecycleHooks.getTaskCompleteHook(this.task.id); + + if (globalCompleteHooks.length === 0 && !taskCompleteHook) { + return; + } + + return await runTimelineMetrics.measureMetric("trigger.dev/execution", "complete", async () => { + for (const hook of globalCompleteHooks) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + hook.name ? `onComplete/${hook.name}` : "onComplete/global", + async (span) => { + await hook.fn({ + payload, + result, + ctx, + signal, + task: this.task.id, + init: initOutput, + }); }, - error: handleErrorResult.error, - }; + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onComplete", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); + + if (hookError) { + throw hookError; } + } - if (handleErrorResult.retry && typeof handleErrorResult.retry === "object") { - const delay = calculateNextRetryDelay(handleErrorResult.retry, execution.attempt.number); + if (taskCompleteHook) { + const [hookError] = await tryCatch( + this._tracer.startActiveSpan( + "onComplete/task", + async (span) => { + await taskCompleteHook({ + payload, + result, + ctx, + signal, + task: this.task.id, + init: initOutput, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onComplete", + [SemanticInternalAttributes.COLLAPSED]: true, + }, + } + ) + ); - return typeof delay === "undefined" - ? { status: "noop", error: handleErrorResult.error } - : { - status: "retry", - retry: { timestamp: Date.now() + delay, delay }, - error: handleErrorResult.error, - }; + if (hookError) { + throw hookError; } + } + }); + } - return { status: "noop", error: handleErrorResult.error }; + #internalErrorResult(execution: TaskRunExecution, code: TaskRunErrorCodes, error: unknown) { + return { + ok: false, + id: execution.run.id, + error: { + type: "INTERNAL_ERROR", + code, + message: + error instanceof Error + ? `${error.name}: ${error.message}` + : typeof error === "string" + ? error + : undefined, + stackTrace: error instanceof Error ? error.stack : undefined, }, - { - attributes: { - [SemanticInternalAttributes.STYLE_ICON]: "exclamation-circle", - }, - } - ); + } satisfies TaskRunExecutionResult; } } diff --git a/packages/core/test/taskExecutor.test.ts b/packages/core/test/taskExecutor.test.ts new file mode 100644 index 0000000000..355471297c --- /dev/null +++ b/packages/core/test/taskExecutor.test.ts @@ -0,0 +1,1716 @@ +import { describe, expect, test } from "vitest"; +import { ConsoleInterceptor } from "../src/v3/consoleInterceptor.js"; +import { + RunFnParams, + ServerBackgroundWorker, + TaskMetadataWithFunctions, + TaskRunErrorCodes, + TaskRunExecution, +} from "../src/v3/index.js"; +import { TracingSDK } from "../src/v3/otel/tracingSDK.js"; +import { TriggerTracer } from "../src/v3/tracer.js"; +import { TaskExecutor } from "../src/v3/workers/taskExecutor.js"; +import { StandardLifecycleHooksManager } from "../src/v3/lifecycleHooks/manager.js"; +import { lifecycleHooks } from "../src/v3/index.js"; + +describe("TaskExecutor", () => { + beforeEach(() => { + lifecycleHooks.setGlobalLifecycleHooksManager(new StandardLifecycleHooksManager()); + }); + + afterEach(() => { + lifecycleHooks.disable(); + }); + + test("should call onComplete with success result", async () => { + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + return { + foo: "bar", + }; + }, + }); + + lifecycleHooks.registerTaskInitHook("test-task", { + id: "test-init", + fn: async () => { + return { + bar: "baz", + }; + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + return { + output: "test-output", + init: params.init, + }; + }, + }, + }; + + const result = await executeTask(task, {}, undefined); + + expect(result).toEqual({ + result: { + ok: true, + id: "test-run-id", + output: '{"json":{"output":"test-output","init":{"foo":"bar","bar":"baz"}}}', + outputType: "application/super+json", + }, + }); + }); + + test("should call onSuccess hooks in correct order with proper data", async () => { + const globalSuccessOrder: string[] = []; + const successPayloads: any[] = []; + const successOutputs: any[] = []; + const successInits: any[] = []; + + // Register global init hook to provide init data + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + return { + foo: "bar", + }; + }, + }); + + // Register two global success hooks + lifecycleHooks.registerGlobalSuccessHook({ + id: "global-success-2", // Register second hook first + fn: async ({ payload, output, init }) => { + console.log("Executing global success hook 2"); + globalSuccessOrder.push("global-2"); + successPayloads.push(payload); + successOutputs.push(output); + successInits.push(init); + }, + }); + + lifecycleHooks.registerGlobalSuccessHook({ + id: "global-success-1", // Register first hook second + fn: async ({ payload, output, init }) => { + console.log("Executing global success hook 1"); + globalSuccessOrder.push("global-1"); + successPayloads.push(payload); + successOutputs.push(output); + successInits.push(init); + }, + }); + + // Register task-specific success hook + lifecycleHooks.registerTaskSuccessHook("test-task", { + id: "task-success", + fn: async ({ payload, output, init }) => { + console.log("Executing task success hook"); + globalSuccessOrder.push("task"); + successPayloads.push(payload); + successOutputs.push(output); + successInits.push(init); + }, + }); + + // Verify hooks are registered + const globalHooks = lifecycleHooks.getGlobalSuccessHooks(); + console.log( + "Registered global hooks:", + globalHooks.map((h) => h.id) + ); + const taskHook = lifecycleHooks.getTaskSuccessHook("test-task"); + console.log("Registered task hook:", taskHook ? "yes" : "no"); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + return { + output: "test-output", + init: params.init, + }; + }, + }, + }; + + const result = await executeTask(task, {}, undefined); + + // Verify hooks were called in correct order - should match registration order + expect(globalSuccessOrder).toEqual(["global-2", "global-1", "task"]); + + // Verify each hook received the correct payload + successPayloads.forEach((payload) => { + expect(payload).toEqual({}); + }); + + // Verify each hook received the correct output + successOutputs.forEach((output) => { + expect(output).toEqual({ + output: "test-output", + init: { foo: "bar" }, + }); + }); + + // Verify each hook received the correct init data + successInits.forEach((init) => { + expect(init).toEqual({ foo: "bar" }); + }); + + // Verify the final result + expect(result).toEqual({ + result: { + ok: true, + id: "test-run-id", + output: '{"json":{"output":"test-output","init":{"foo":"bar"}}}', + outputType: "application/super+json", + }, + }); + }); + + test("should call onStart hooks in correct order with proper data", async () => { + const globalStartOrder: string[] = []; + const startPayloads: any[] = []; + const startInits: any[] = []; + + // Register global init hook to provide init data + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + return { + foo: "bar", + }; + }, + }); + + // Register two global start hooks + lifecycleHooks.registerGlobalStartHook({ + id: "global-start-1", + fn: async ({ payload, ctx, init }) => { + console.log("Executing global start hook 1"); + globalStartOrder.push("global-1"); + startPayloads.push(payload); + startInits.push(init); + }, + }); + + lifecycleHooks.registerGlobalStartHook({ + id: "global-start-2", + fn: async ({ payload, ctx, init }) => { + console.log("Executing global start hook 2"); + globalStartOrder.push("global-2"); + startPayloads.push(payload); + startInits.push(init); + }, + }); + + // Register task-specific start hook + lifecycleHooks.registerTaskStartHook("test-task", { + id: "task-start", + fn: async ({ payload, ctx, init }) => { + console.log("Executing task start hook"); + globalStartOrder.push("task"); + startPayloads.push(payload); + startInits.push(init); + }, + }); + + // Verify hooks are registered + const globalHooks = lifecycleHooks.getGlobalStartHooks(); + console.log( + "Registered global hooks:", + globalHooks.map((h) => h.id) + ); + const taskHook = lifecycleHooks.getTaskStartHook("test-task"); + console.log("Registered task hook:", taskHook ? "yes" : "no"); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + return { + output: "test-output", + init: params.init, + }; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify hooks were called in correct order + expect(globalStartOrder).toEqual(["global-1", "global-2", "task"]); + + // Verify each hook received the correct payload + startPayloads.forEach((payload) => { + expect(payload).toEqual({ test: "data" }); + }); + + console.log("startInits", startInits); + + // Verify each hook received the correct init data + startInits.forEach((init) => { + expect(init).toEqual({ foo: "bar" }); + }); + + // Verify the final result + expect(result).toEqual({ + result: { + ok: true, + id: "test-run-id", + output: '{"json":{"output":"test-output","init":{"foo":"bar"}}}', + outputType: "application/super+json", + }, + }); + }); + + test("should call onFailure hooks with error when task fails", async () => { + const globalFailureOrder: string[] = []; + const failurePayloads: any[] = []; + const failureErrors: any[] = []; + const failureInits: any[] = []; + + // Register global init hook to provide init data + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + return { + foo: "bar", + }; + }, + }); + + // Register two global failure hooks + lifecycleHooks.registerGlobalFailureHook({ + id: "global-failure-1", + fn: async ({ payload, error, init }) => { + console.log("Executing global failure hook 1"); + globalFailureOrder.push("global-1"); + failurePayloads.push(payload); + failureErrors.push(error); + failureInits.push(init); + }, + }); + + lifecycleHooks.registerGlobalFailureHook({ + id: "global-failure-2", + fn: async ({ payload, error, init }) => { + console.log("Executing global failure hook 2"); + globalFailureOrder.push("global-2"); + failurePayloads.push(payload); + failureErrors.push(error); + failureInits.push(init); + }, + }); + + // Register task-specific failure hook + lifecycleHooks.registerTaskFailureHook("test-task", { + id: "task-failure", + fn: async ({ payload, error, init }) => { + console.log("Executing task failure hook"); + globalFailureOrder.push("task"); + failurePayloads.push(payload); + failureErrors.push(error); + failureInits.push(init); + }, + }); + + // Verify hooks are registered + const globalHooks = lifecycleHooks.getGlobalFailureHooks(); + console.log( + "Registered global hooks:", + globalHooks.map((h) => h.id) + ); + const taskHook = lifecycleHooks.getTaskFailureHook("test-task"); + console.log("Registered task hook:", taskHook ? "yes" : "no"); + + const expectedError = new Error("Task failed intentionally"); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + throw expectedError; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify hooks were called in correct order + expect(globalFailureOrder).toEqual(["global-1", "global-2", "task"]); + + // Verify each hook received the correct payload + failurePayloads.forEach((payload) => { + expect(payload).toEqual({ test: "data" }); + }); + + // Verify each hook received the correct error + failureErrors.forEach((error) => { + expect(error).toBe(expectedError); + }); + + // Verify each hook received the correct init data + failureInits.forEach((init) => { + expect(init).toEqual({ foo: "bar" }); + }); + + // Verify the final result contains the error + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "BUILT_IN_ERROR", + message: "Task failed intentionally", + name: "Error", + stackTrace: expect.any(String), + }, + skippedRetrying: false, + }, + }); + }); + + test("should call onComplete hooks in correct order with proper data", async () => { + const globalCompleteOrder: string[] = []; + const completePayloads: any[] = []; + const completeResults: any[] = []; + const completeInits: any[] = []; + + // Register global init hook to provide init data + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + return { + foo: "bar", + }; + }, + }); + + // Register two global complete hooks + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete-1", + fn: async ({ payload, result, init }) => { + console.log("Executing global complete hook 1"); + globalCompleteOrder.push("global-1"); + completePayloads.push(payload); + completeResults.push(result); + completeInits.push(init); + }, + }); + + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete-2", + fn: async ({ payload, result, init }) => { + console.log("Executing global complete hook 2"); + globalCompleteOrder.push("global-2"); + completePayloads.push(payload); + completeResults.push(result); + completeInits.push(init); + }, + }); + + // Register task-specific complete hook + lifecycleHooks.registerTaskCompleteHook("test-task", { + id: "task-complete", + fn: async ({ payload, result, init }) => { + console.log("Executing task complete hook"); + globalCompleteOrder.push("task"); + completePayloads.push(payload); + completeResults.push(result); + completeInits.push(init); + }, + }); + + // Verify hooks are registered + const globalHooks = lifecycleHooks.getGlobalCompleteHooks(); + console.log( + "Registered global hooks:", + globalHooks.map((h) => h.id) + ); + const taskHook = lifecycleHooks.getTaskCompleteHook("test-task"); + console.log("Registered task hook:", taskHook ? "yes" : "no"); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + return { + output: "test-output", + init: params.init, + }; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify hooks were called in correct order + expect(globalCompleteOrder).toEqual(["global-1", "global-2", "task"]); + + // Verify each hook received the correct payload + completePayloads.forEach((payload) => { + expect(payload).toEqual({ test: "data" }); + }); + + // Verify each hook received the correct result + completeResults.forEach((result) => { + expect(result).toEqual({ + ok: true, + data: { + output: "test-output", + init: { foo: "bar" }, + }, + }); + }); + + // Verify each hook received the correct init data + completeInits.forEach((init) => { + expect(init).toEqual({ foo: "bar" }); + }); + + // Verify the final result + expect(result).toEqual({ + result: { + ok: true, + id: "test-run-id", + output: '{"json":{"output":"test-output","init":{"foo":"bar"}}}', + outputType: "application/super+json", + }, + }); + }); + + test("should call onComplete hooks with error when task fails", async () => { + const globalCompleteOrder: string[] = []; + const completePayloads: any[] = []; + const completeResults: any[] = []; + const completeInits: any[] = []; + + // Register global init hook to provide init data + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + return { + foo: "bar", + }; + }, + }); + + // Register global complete hooks + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete", + fn: async ({ payload, result, init }) => { + console.log("Executing global complete hook"); + globalCompleteOrder.push("global"); + completePayloads.push(payload); + completeResults.push(result); + completeInits.push(init); + }, + }); + + // Register task-specific complete hook + lifecycleHooks.registerTaskCompleteHook("test-task", { + id: "task-complete", + fn: async ({ payload, result, init }) => { + console.log("Executing task complete hook"); + globalCompleteOrder.push("task"); + completePayloads.push(payload); + completeResults.push(result); + completeInits.push(init); + }, + }); + + const expectedError = new Error("Task failed intentionally"); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + throw expectedError; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify hooks were called in correct order + expect(globalCompleteOrder).toEqual(["global", "task"]); + + // Verify each hook received the correct payload + completePayloads.forEach((payload) => { + expect(payload).toEqual({ test: "data" }); + }); + + // Verify each hook received the error result + completeResults.forEach((result) => { + expect(result).toEqual({ + ok: false, + error: expectedError, + }); + }); + + // Verify each hook received the correct init data + completeInits.forEach((init) => { + expect(init).toEqual({ foo: "bar" }); + }); + + // Verify the final result contains the error + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "BUILT_IN_ERROR", + message: "Task failed intentionally", + name: "Error", + stackTrace: expect.any(String), + }, + skippedRetrying: false, + }, + }); + }); + + test("should call catchError hooks in correct order and stop at first handler that returns a result", async () => { + const hookCallOrder: string[] = []; + const expectedError = new Error("Task failed intentionally"); + + // Register global init hook to provide init data + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + return { + foo: "bar", + }; + }, + }); + + // Register task-specific catch error hook that doesn't handle the error + lifecycleHooks.registerTaskCatchErrorHook("test-task", { + id: "task-catch-error", + fn: async ({ payload, error, init, retry }) => { + console.log("Executing task catch error hook"); + hookCallOrder.push("task"); + // Return undefined to let it fall through to global handlers + return undefined; + }, + }); + + // Register first global catch error hook that doesn't handle the error + lifecycleHooks.registerGlobalCatchErrorHook({ + id: "global-catch-error-1", + fn: async ({ payload, error, init, retry }) => { + console.log("Executing global catch error hook 1"); + hookCallOrder.push("global-1"); + // Return undefined to let it fall through to next handler + return undefined; + }, + }); + + // Register second global catch error hook that handles the error + lifecycleHooks.registerGlobalCatchErrorHook({ + id: "global-catch-error-2", + fn: async ({ payload, error, init, retry }) => { + console.log("Executing global catch error hook 2"); + hookCallOrder.push("global-2"); + // Return a result to handle the error + return { + retry: { + maxAttempts: 3, + minDelay: 1000, + maxDelay: 5000, + factor: 2, + }, + }; + }, + }); + + // Register third global catch error hook that should never be called + lifecycleHooks.registerGlobalCatchErrorHook({ + id: "global-catch-error-3", + fn: async ({ payload, error, init, retry }) => { + console.log("Executing global catch error hook 3"); + hookCallOrder.push("global-3"); + return undefined; + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + throw expectedError; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify hooks were called in correct order and stopped after second global hook + expect(hookCallOrder).toEqual(["task", "global-1", "global-2"]); + // global-3 should not be called since global-2 returned a result + + // Verify the final result contains retry information from the second global hook + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "BUILT_IN_ERROR", + message: "Task failed intentionally", + name: "Error", + stackTrace: expect.any(String), + }, + retry: { + timestamp: expect.any(Number), + delay: expect.any(Number), + }, + skippedRetrying: false, + }, + }); + }); + + test("should skip retrying if catch error hook returns skipRetrying", async () => { + const hookCallOrder: string[] = []; + const expectedError = new Error("Task failed intentionally"); + + // Register task-specific catch error hook that handles the error + lifecycleHooks.registerTaskCatchErrorHook("test-task", { + id: "task-catch-error", + fn: async ({ payload, error, init }) => { + console.log("Executing task catch error hook"); + hookCallOrder.push("task"); + return { + skipRetrying: true, + error: new Error("Modified error in catch hook"), + }; + }, + }); + + // Register global catch error hook that should never be called + lifecycleHooks.registerGlobalCatchErrorHook({ + id: "global-catch-error", + fn: async ({ payload, error, init }) => { + console.log("Executing global catch error hook"); + hookCallOrder.push("global"); + return undefined; + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + throw expectedError; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify only task hook was called + expect(hookCallOrder).toEqual(["task"]); + + // Verify the final result shows skipped retrying and the modified error + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "BUILT_IN_ERROR", + message: "Modified error in catch hook", + name: "Error", + stackTrace: expect.any(String), + }, + skippedRetrying: true, + }, + }); + }); + + test("should use specific retry timing if catch error hook provides it", async () => { + const hookCallOrder: string[] = []; + const expectedError = new Error("Task failed intentionally"); + const specificRetryDate = new Date(Date.now() + 30000); // 30 seconds in future + + // Register task-specific catch error hook that specifies retry timing + lifecycleHooks.registerTaskCatchErrorHook("test-task", { + id: "task-catch-error", + fn: async ({ payload, error, init }) => { + console.log("Executing task catch error hook"); + hookCallOrder.push("task"); + return { + retryAt: specificRetryDate, + error: expectedError, + }; + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + throw expectedError; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify only task hook was called + expect(hookCallOrder).toEqual(["task"]); + + // Verify the final result contains the specific retry timing + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "BUILT_IN_ERROR", + message: "Task failed intentionally", + name: "Error", + stackTrace: expect.any(String), + }, + retry: { + timestamp: specificRetryDate.getTime(), + delay: expect.any(Number), + }, + skippedRetrying: false, + }, + }); + + expect((result as any).result.retry.delay).toBeGreaterThan(29900); + expect((result as any).result.retry.delay).toBeLessThan(30100); + }); + + test("should execute middleware hooks in correct order around other hooks", async () => { + const executionOrder: string[] = []; + + // Register global init hook + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + executionOrder.push("init"); + return { + foo: "bar", + }; + }, + }); + + // Register global start hook + lifecycleHooks.registerGlobalStartHook({ + id: "global-start", + fn: async ({ payload }) => { + executionOrder.push("start"); + }, + }); + + // Register global success hook + lifecycleHooks.registerGlobalSuccessHook({ + id: "global-success", + fn: async ({ payload, output }) => { + executionOrder.push("success"); + }, + }); + + // Register global complete hook + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete", + fn: async ({ payload, result }) => { + executionOrder.push("complete"); + }, + }); + + // Register task-specific middleware that executes first + lifecycleHooks.registerTaskMiddlewareHook("test-task", { + id: "task-middleware", + fn: async ({ payload, ctx, next }) => { + executionOrder.push("task-middleware-before"); + await next(); + executionOrder.push("task-middleware-after"); + }, + }); + + // Register two global middleware hooks + lifecycleHooks.registerGlobalMiddlewareHook({ + id: "global-middleware-1", + fn: async ({ payload, ctx, next }) => { + executionOrder.push("global-middleware-1-before"); + await next(); + executionOrder.push("global-middleware-1-after"); + }, + }); + + lifecycleHooks.registerGlobalMiddlewareHook({ + id: "global-middleware-2", + fn: async ({ payload, ctx, next }) => { + executionOrder.push("global-middleware-2-before"); + await next(); + executionOrder.push("global-middleware-2-after"); + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + executionOrder.push("run"); + return { + output: "test-output", + init: params.init, + }; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify the execution order: + // 1. Global middlewares (outside to inside) + // 2. Task middleware + // 3. Init hook + // 4. Start hook + // 5. Run function + // 6. Success hook + // 7. Complete hook + // 8. Middlewares in reverse order + expect(executionOrder).toEqual([ + "global-middleware-1-before", + "global-middleware-2-before", + "task-middleware-before", + "init", + "start", + "run", + "success", + "complete", + "task-middleware-after", + "global-middleware-2-after", + "global-middleware-1-after", + ]); + + // Verify the final result + expect(result).toEqual({ + result: { + ok: true, + id: "test-run-id", + output: '{"json":{"output":"test-output","init":{"foo":"bar"}}}', + outputType: "application/super+json", + }, + }); + }); + + test("should handle middleware errors correctly", async () => { + const executionOrder: string[] = []; + const expectedError = new Error("Middleware error"); + + // Register global middleware that throws an error + lifecycleHooks.registerGlobalMiddlewareHook({ + id: "global-middleware", + fn: async ({ payload, ctx, next }) => { + executionOrder.push("middleware-before"); + throw expectedError; + // Should never get here + await next(); + executionOrder.push("middleware-after"); + }, + }); + + // Register failure hook to verify it's called + lifecycleHooks.registerGlobalFailureHook({ + id: "global-failure", + fn: async ({ payload, error }) => { + executionOrder.push("failure"); + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + executionOrder.push("run"); + return { + output: "test-output", + }; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify only the middleware-before hook ran + expect(executionOrder).toEqual(["middleware-before"]); + + // Verify the error result + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "INTERNAL_ERROR", + message: "Error: Middleware error", + code: "TASK_MIDDLEWARE_ERROR", + stackTrace: expect.any(String), + }, + }, + }); + }); + + test("should propagate errors from init hooks", async () => { + const executionOrder: string[] = []; + const expectedError = new Error("Init hook error"); + + // Register global init hook that throws an error + lifecycleHooks.registerGlobalInitHook({ + id: "failing-init", + fn: async () => { + executionOrder.push("global-init"); + throw expectedError; + }, + }); + + // Register task init hook that should never be called + lifecycleHooks.registerTaskInitHook("test-task", { + id: "task-init", + fn: async () => { + executionOrder.push("task-init"); + return { + foo: "bar", + }; + }, + }); + + // Register failure hook to verify it's called + lifecycleHooks.registerGlobalFailureHook({ + id: "global-failure", + fn: async ({ error }) => { + executionOrder.push("failure"); + expect(error).toBe(expectedError); + }, + }); + + // Register complete hook to verify it's called with error + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete", + fn: async ({ result }) => { + executionOrder.push("complete"); + expect(result).toEqual({ + ok: false, + error: expectedError, + }); + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + executionOrder.push("run"); + return { + output: "test-output", + }; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify only the global init hook ran, and failure/complete hooks were called + expect(executionOrder).toEqual(["global-init", "failure", "complete"]); + + // Verify the error result + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "BUILT_IN_ERROR", + message: "Init hook error", + name: "Error", + stackTrace: expect.any(String), + }, + skippedRetrying: false, + }, + }); + }); + + test("should propagate errors from task init hooks", async () => { + const executionOrder: string[] = []; + const expectedError = new Error("Task init hook error"); + + // Register global init hook that succeeds + lifecycleHooks.registerGlobalInitHook({ + id: "global-init", + fn: async () => { + executionOrder.push("global-init"); + return { + foo: "bar", + }; + }, + }); + + // Register task init hook that throws an error + lifecycleHooks.registerTaskInitHook("test-task", { + id: "task-init", + fn: async () => { + executionOrder.push("task-init"); + throw expectedError; + }, + }); + + // Register failure hook to verify it's called + lifecycleHooks.registerGlobalFailureHook({ + id: "global-failure", + fn: async ({ error, init }) => { + executionOrder.push("failure"); + expect(error).toBe(expectedError); + }, + }); + + // Register complete hook to verify it's called with error + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete", + fn: async ({ result, init }) => { + executionOrder.push("complete"); + expect(result).toEqual({ + ok: false, + error: expectedError, + }); + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + executionOrder.push("run"); + return { + output: "test-output", + }; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify both init hooks ran, but run wasn't called, and failure/complete hooks were called + expect(executionOrder).toEqual(["global-init", "task-init", "failure", "complete"]); + + // Verify the error result + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "BUILT_IN_ERROR", + message: "Task init hook error", + name: "Error", + stackTrace: expect.any(String), + }, + skippedRetrying: false, + }, + }); + }); + + test("should propagate errors from start hooks", async () => { + const executionOrder: string[] = []; + const expectedError = new Error("Start hook error"); + + // Register global init hook that succeeds + lifecycleHooks.registerGlobalInitHook({ + id: "global-init", + fn: async () => { + executionOrder.push("global-init"); + return { + foo: "bar", + }; + }, + }); + + // Register global start hook that throws an error + lifecycleHooks.registerGlobalStartHook({ + id: "global-start", + fn: async () => { + executionOrder.push("global-start"); + throw expectedError; + }, + }); + + // Register task start hook that should never be called + lifecycleHooks.registerTaskStartHook("test-task", { + id: "task-start", + fn: async () => { + executionOrder.push("task-start"); + }, + }); + + // Register failure hook to verify it's called + lifecycleHooks.registerGlobalFailureHook({ + id: "global-failure", + fn: async ({ error, init }) => { + executionOrder.push("failure"); + expect(error).toBe(expectedError); + // Verify we got the init data + expect(init).toEqual({ foo: "bar" }); + }, + }); + + // Register complete hook to verify it's called with error + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete", + fn: async ({ result, init }) => { + executionOrder.push("complete"); + expect(result).toEqual({ + ok: false, + error: expectedError, + }); + // Verify we got the init data + expect(init).toEqual({ foo: "bar" }); + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + executionOrder.push("run"); + return { + output: "test-output", + }; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify init succeeded, start hook failed, and run wasn't called + expect(executionOrder).toEqual(["global-init", "global-start", "failure", "complete"]); + + // Verify the error result + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "BUILT_IN_ERROR", + message: "Start hook error", + name: "Error", + stackTrace: expect.any(String), + }, + skippedRetrying: false, + }, + }); + }); + + test("should call cleanup hooks in correct order after other hooks but before middleware completion", async () => { + const executionOrder: string[] = []; + + // Register global init hook + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + executionOrder.push("init"); + return { + foo: "bar", + }; + }, + }); + + // Register global start hook + lifecycleHooks.registerGlobalStartHook({ + id: "global-start", + fn: async () => { + executionOrder.push("start"); + }, + }); + + // Register global success hook + lifecycleHooks.registerGlobalSuccessHook({ + id: "global-success", + fn: async () => { + executionOrder.push("success"); + }, + }); + + // Register global complete hook + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete", + fn: async () => { + executionOrder.push("complete"); + }, + }); + + // Register global cleanup hooks + lifecycleHooks.registerGlobalCleanupHook({ + id: "global-cleanup-1", + fn: async ({ init }) => { + executionOrder.push("global-cleanup-1"); + // Verify we have access to init data + expect(init).toEqual({ foo: "bar" }); + }, + }); + + lifecycleHooks.registerGlobalCleanupHook({ + id: "global-cleanup-2", + fn: async ({ init }) => { + executionOrder.push("global-cleanup-2"); + // Verify we have access to init data + expect(init).toEqual({ foo: "bar" }); + }, + }); + + // Register task-specific cleanup hook + lifecycleHooks.registerTaskCleanupHook("test-task", { + id: "task-cleanup", + fn: async ({ init }) => { + executionOrder.push("task-cleanup"); + // Verify we have access to init data + expect(init).toEqual({ foo: "bar" }); + }, + }); + + // Register middleware to verify cleanup happens before middleware completion + lifecycleHooks.registerGlobalMiddlewareHook({ + id: "global-middleware", + fn: async ({ next }) => { + executionOrder.push("middleware-before"); + await next(); + executionOrder.push("middleware-after"); + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + executionOrder.push("run"); + return { + output: "test-output", + init: params.init, + }; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify the execution order: + // 1. Middleware starts + // 2. Init hook + // 3. Start hook + // 4. Run function + // 5. Success hook + // 6. Complete hook + // 7. Cleanup hooks + // 8. Middleware completes + expect(executionOrder).toEqual([ + "middleware-before", + "init", + "start", + "run", + "success", + "complete", + "global-cleanup-1", + "global-cleanup-2", + "task-cleanup", + "middleware-after", + ]); + + // Verify the final result + expect(result).toEqual({ + result: { + ok: true, + id: "test-run-id", + output: '{"json":{"output":"test-output","init":{"foo":"bar"}}}', + outputType: "application/super+json", + }, + }); + }); + + test("should call cleanup hooks even when task fails", async () => { + const executionOrder: string[] = []; + const expectedError = new Error("Task failed intentionally"); + + // Register global init hook + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + executionOrder.push("init"); + return { + foo: "bar", + }; + }, + }); + + // Register failure hook + lifecycleHooks.registerGlobalFailureHook({ + id: "global-failure", + fn: async () => { + executionOrder.push("failure"); + }, + }); + + // Register complete hook + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete", + fn: async () => { + executionOrder.push("complete"); + }, + }); + + // Register cleanup hooks + lifecycleHooks.registerGlobalCleanupHook({ + id: "global-cleanup", + fn: async ({ init }) => { + executionOrder.push("global-cleanup"); + // Verify we have access to init data even after failure + expect(init).toEqual({ foo: "bar" }); + }, + }); + + lifecycleHooks.registerTaskCleanupHook("test-task", { + id: "task-cleanup", + fn: async ({ init }) => { + executionOrder.push("task-cleanup"); + // Verify we have access to init data even after failure + expect(init).toEqual({ foo: "bar" }); + }, + }); + + const task = { + id: "test-task", + fns: { + run: async () => { + executionOrder.push("run"); + throw expectedError; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, undefined); + + // Verify cleanup hooks are called even after failure + expect(executionOrder).toEqual([ + "init", + "run", + "failure", + "complete", + "global-cleanup", + "task-cleanup", + ]); + + // Verify the error result + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "BUILT_IN_ERROR", + message: "Task failed intentionally", + name: "Error", + stackTrace: expect.any(String), + }, + skippedRetrying: false, + }, + }); + }); + + test("should handle max duration abort signal and call hooks in correct order", async () => { + const executionOrder: string[] = []; + const maxDurationMs = 1000; + + // Create an abort controller that we'll trigger manually + const controller = new AbortController(); + + // Register global init hook + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + executionOrder.push("init"); + return { + foo: "bar", + }; + }, + }); + + // Register failure hook + lifecycleHooks.registerGlobalFailureHook({ + id: "global-failure", + fn: async ({ error }) => { + executionOrder.push("failure"); + expect((error as Error).message).toBe( + `Task execution exceeded maximum duration of ${maxDurationMs}ms` + ); + }, + }); + + // Register complete hook + lifecycleHooks.registerGlobalCompleteHook({ + id: "global-complete", + fn: async ({ result }) => { + executionOrder.push("complete"); + expect(result.ok).toBe(false); + }, + }); + + // Register cleanup hook + lifecycleHooks.registerGlobalCleanupHook({ + id: "global-cleanup", + fn: async () => { + executionOrder.push("cleanup"); + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + executionOrder.push("run-start"); + + // Create a promise that never resolves + await new Promise((resolve) => { + // Trigger abort after a small delay + setTimeout(() => { + controller.abort(); + }, 10); + }); + + // This should never be reached + executionOrder.push("run-end"); + }, + }, + }; + + const result = await executeTask(task, { test: "data" }, controller.signal); + + // Verify hooks were called in correct order + expect(executionOrder).toEqual(["init", "run-start", "failure", "complete", "cleanup"]); + + // Verify the error result + expect(result).toEqual({ + result: { + ok: false, + id: "test-run-id", + error: { + type: "INTERNAL_ERROR", + code: TaskRunErrorCodes.MAX_DURATION_EXCEEDED, + message: "Task execution exceeded maximum duration of 1000ms", + stackTrace: expect.any(String), + }, + skippedRetrying: false, + }, + }); + }); + + test("should call onWait and onResume hooks in correct order with proper data", async () => { + const executionOrder: string[] = []; + const waitData = { type: "task", runId: "test-run-id" } as const; + + // Register global init hook to provide init data + lifecycleHooks.registerGlobalInitHook({ + id: "test-init", + fn: async () => { + executionOrder.push("init"); + return { + foo: "bar", + }; + }, + }); + + // Register global wait hooks + lifecycleHooks.registerGlobalWaitHook({ + id: "global-wait-1", + fn: async ({ payload, wait, init }) => { + executionOrder.push("global-wait-1"); + expect(wait).toEqual(waitData); + expect(init).toEqual({ foo: "bar" }); + }, + }); + + lifecycleHooks.registerGlobalWaitHook({ + id: "global-wait-2", + fn: async ({ payload, wait, init }) => { + executionOrder.push("global-wait-2"); + expect(wait).toEqual(waitData); + expect(init).toEqual({ foo: "bar" }); + }, + }); + + // Register task-specific wait hook + lifecycleHooks.registerTaskWaitHook("test-task", { + id: "task-wait", + fn: async ({ payload, wait, init }) => { + executionOrder.push("task-wait"); + expect(wait).toEqual(waitData); + expect(init).toEqual({ foo: "bar" }); + }, + }); + + // Register global resume hooks + lifecycleHooks.registerGlobalResumeHook({ + id: "global-resume-1", + fn: async ({ payload, wait, init }) => { + executionOrder.push("global-resume-1"); + expect(wait).toEqual(waitData); + expect(init).toEqual({ foo: "bar" }); + }, + }); + + lifecycleHooks.registerGlobalResumeHook({ + id: "global-resume-2", + fn: async ({ payload, wait, init }) => { + executionOrder.push("global-resume-2"); + expect(wait).toEqual(waitData); + expect(init).toEqual({ foo: "bar" }); + }, + }); + + // Register task-specific resume hook + lifecycleHooks.registerTaskResumeHook("test-task", { + id: "task-resume", + fn: async ({ payload, wait, init }) => { + executionOrder.push("task-resume"); + expect(wait).toEqual(waitData); + expect(init).toEqual({ foo: "bar" }); + }, + }); + + const task = { + id: "test-task", + fns: { + run: async (payload: any, params: RunFnParams) => { + executionOrder.push("run-start"); + + // Simulate a wait + await lifecycleHooks.callOnWaitHookListeners(waitData); + + // Simulate some time passing + await new Promise((resolve) => setTimeout(resolve, 10)); + + // Simulate resuming + await lifecycleHooks.callOnResumeHookListeners(waitData); + + executionOrder.push("run-end"); + return { success: true }; + }, + }, + }; + + const result = await executeTask(task, { test: "data" }); + + // Verify hooks were called in correct order + expect(executionOrder).toEqual([ + "init", + "run-start", + "global-wait-1", + "global-wait-2", + "task-wait", + "global-resume-1", + "global-resume-2", + "task-resume", + "run-end", + ]); + + // Verify the final result + expect(result).toEqual({ + result: { + ok: true, + id: "test-run-id", + output: '{"json":{"success":true}}', + outputType: "application/super+json", + }, + }); + }); +}); + +function executeTask(task: TaskMetadataWithFunctions, payload: any, signal?: AbortSignal) { + const tracingSDK = new TracingSDK({ + url: "http://localhost:4318", + }); + + const tracer = new TriggerTracer({ + name: "test-task", + version: "1.0.0", + tracer: tracingSDK.getTracer("test-task"), + logger: tracingSDK.getLogger("test-task"), + }); + + const consoleInterceptor = new ConsoleInterceptor(tracingSDK.getLogger("test-task"), false); + + const executor = new TaskExecutor(task, { + tracingSDK, + tracer, + consoleInterceptor, + retries: { + enabledInDev: false, + default: { + maxAttempts: 1, + }, + }, + }); + + const execution: TaskRunExecution = { + task: { + id: "test-task", + filePath: "test-task.ts", + }, + attempt: { + number: 1, + startedAt: new Date(), + id: "test-attempt-id", + status: "success", + backgroundWorkerId: "test-background-worker-id", + backgroundWorkerTaskId: "test-background-worker-task-id", + }, + run: { + id: "test-run-id", + payload: JSON.stringify(payload), + payloadType: "application/json", + metadata: {}, + startedAt: new Date(), + tags: [], + isTest: false, + createdAt: new Date(), + durationMs: 0, + costInCents: 0, + baseCostInCents: 0, + priority: 0, + maxDuration: 1000, + }, + machine: { + name: "micro", + cpu: 1, + memory: 1, + centsPerMs: 0, + }, + queue: { + name: "test-queue", + id: "test-queue-id", + }, + environment: { + type: "PRODUCTION", + id: "test-environment-id", + slug: "test-environment-slug", + }, + organization: { + id: "test-organization-id", + name: "test-organization-name", + slug: "test-organization-slug", + }, + project: { + id: "test-project-id", + name: "test-project-name", + slug: "test-project-slug", + ref: "test-project-ref", + }, + }; + + const worker: ServerBackgroundWorker = { + id: "test-background-worker-id", + version: "1.0.0", + contentHash: "test-content-hash", + engine: "V2", + }; + + return executor.execute(execution, worker, {}, signal); +} diff --git a/packages/trigger-sdk/src/v3/hooks.ts b/packages/trigger-sdk/src/v3/hooks.ts new file mode 100644 index 0000000000..d864f4ec8e --- /dev/null +++ b/packages/trigger-sdk/src/v3/hooks.ts @@ -0,0 +1,133 @@ +import { + lifecycleHooks, + type AnyOnStartHookFunction, + type TaskStartHookParams, + type OnStartHookFunction, + type AnyOnFailureHookFunction, + type AnyOnSuccessHookFunction, + type AnyOnCompleteHookFunction, + type TaskCompleteResult, + type AnyOnWaitHookFunction, + type AnyOnResumeHookFunction, + type AnyOnCatchErrorHookFunction, + type AnyOnMiddlewareHookFunction, +} from "@trigger.dev/core/v3"; + +export type { + AnyOnStartHookFunction, + TaskStartHookParams, + OnStartHookFunction, + AnyOnFailureHookFunction, + AnyOnSuccessHookFunction, + AnyOnCompleteHookFunction, + TaskCompleteResult, + AnyOnWaitHookFunction, + AnyOnResumeHookFunction, + AnyOnCatchErrorHookFunction, + AnyOnMiddlewareHookFunction, +}; + +export function onStart(name: string, fn: AnyOnStartHookFunction): void; +export function onStart(fn: AnyOnStartHookFunction): void; +export function onStart( + fnOrName: string | AnyOnStartHookFunction, + fn?: AnyOnStartHookFunction +): void { + lifecycleHooks.registerGlobalStartHook({ + id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined, + fn: typeof fnOrName === "function" ? fnOrName : fn!, + }); +} + +export function onFailure(name: string, fn: AnyOnFailureHookFunction): void; +export function onFailure(fn: AnyOnFailureHookFunction): void; +export function onFailure( + fnOrName: string | AnyOnFailureHookFunction, + fn?: AnyOnFailureHookFunction +): void { + lifecycleHooks.registerGlobalFailureHook({ + id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined, + fn: typeof fnOrName === "function" ? fnOrName : fn!, + }); +} + +export function onSuccess(name: string, fn: AnyOnSuccessHookFunction): void; +export function onSuccess(fn: AnyOnSuccessHookFunction): void; +export function onSuccess( + fnOrName: string | AnyOnSuccessHookFunction, + fn?: AnyOnSuccessHookFunction +): void { + lifecycleHooks.registerGlobalSuccessHook({ + id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined, + fn: typeof fnOrName === "function" ? fnOrName : fn!, + }); +} + +export function onComplete(name: string, fn: AnyOnCompleteHookFunction): void; +export function onComplete(fn: AnyOnCompleteHookFunction): void; +export function onComplete( + fnOrName: string | AnyOnCompleteHookFunction, + fn?: AnyOnCompleteHookFunction +): void { + lifecycleHooks.registerGlobalCompleteHook({ + id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined, + fn: typeof fnOrName === "function" ? fnOrName : fn!, + }); +} + +export function onWait(name: string, fn: AnyOnWaitHookFunction): void; +export function onWait(fn: AnyOnWaitHookFunction): void; +export function onWait(fnOrName: string | AnyOnWaitHookFunction, fn?: AnyOnWaitHookFunction): void { + lifecycleHooks.registerGlobalWaitHook({ + id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined, + fn: typeof fnOrName === "function" ? fnOrName : fn!, + }); +} + +export function onResume(name: string, fn: AnyOnResumeHookFunction): void; +export function onResume(fn: AnyOnResumeHookFunction): void; +export function onResume( + fnOrName: string | AnyOnResumeHookFunction, + fn?: AnyOnResumeHookFunction +): void { + lifecycleHooks.registerGlobalResumeHook({ + id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined, + fn: typeof fnOrName === "function" ? fnOrName : fn!, + }); +} + +/** @deprecated Use onCatchError instead */ +export function onHandleError(name: string, fn: AnyOnCatchErrorHookFunction): void; +/** @deprecated Use onCatchError instead */ +export function onHandleError(fn: AnyOnCatchErrorHookFunction): void; +/** @deprecated Use onCatchError instead */ +export function onHandleError( + fnOrName: string | AnyOnCatchErrorHookFunction, + fn?: AnyOnCatchErrorHookFunction +): void { + onCatchError(fnOrName as any, fn as any); +} + +export function onCatchError(name: string, fn: AnyOnCatchErrorHookFunction): void; +export function onCatchError(fn: AnyOnCatchErrorHookFunction): void; +export function onCatchError( + fnOrName: string | AnyOnCatchErrorHookFunction, + fn?: AnyOnCatchErrorHookFunction +): void { + lifecycleHooks.registerGlobalCatchErrorHook({ + id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined, + fn: typeof fnOrName === "function" ? fnOrName : fn!, + }); +} + +export function middleware(name: string, fn: AnyOnMiddlewareHookFunction): void; +export function middleware(fn: AnyOnMiddlewareHookFunction): void; +export function middleware( + fnOrName: string | AnyOnMiddlewareHookFunction, + fn?: AnyOnMiddlewareHookFunction +): void { + lifecycleHooks.registerGlobalMiddlewareHook({ + id: typeof fnOrName === "string" ? fnOrName : fnOrName.name ? fnOrName.name : undefined, + fn: typeof fnOrName === "function" ? fnOrName : fn!, + }); +} diff --git a/packages/trigger-sdk/src/v3/index.ts b/packages/trigger-sdk/src/v3/index.ts index 5f00a4a3e1..4837525734 100644 --- a/packages/trigger-sdk/src/v3/index.ts +++ b/packages/trigger-sdk/src/v3/index.ts @@ -12,6 +12,7 @@ export * from "./tags.js"; export * from "./metadata.js"; export * from "./timeout.js"; export * from "./webhooks.js"; +export * from "./locals.js"; export type { Context }; import type { Context } from "./shared.js"; diff --git a/packages/trigger-sdk/src/v3/locals.ts b/packages/trigger-sdk/src/v3/locals.ts new file mode 100644 index 0000000000..1e6082719b --- /dev/null +++ b/packages/trigger-sdk/src/v3/locals.ts @@ -0,0 +1,5 @@ +import { type Locals, locals, type LocalsKey } from "@trigger.dev/core/v3"; + +export type { Locals, LocalsKey }; + +export { locals }; diff --git a/packages/trigger-sdk/src/v3/shared.ts b/packages/trigger-sdk/src/v3/shared.ts index 7e7ad651cb..e21cc2f219 100644 --- a/packages/trigger-sdk/src/v3/shared.ts +++ b/packages/trigger-sdk/src/v3/shared.ts @@ -1,4 +1,4 @@ -import { SpanKind, SpanStatusCode } from "@opentelemetry/api"; +import { SpanKind } from "@opentelemetry/api"; import { SerializableJson } from "@trigger.dev/core"; import { accessoryAttributes, @@ -8,45 +8,59 @@ import { convertToolParametersToSchema, createErrorTaskError, defaultRetryOptions, + flattenIdempotencyKey, + getEnvVar, getSchemaParseFn, InitOutput, + lifecycleHooks, makeIdempotencyKey, parsePacket, Queue, QueueOptions, + resourceCatalog, runtime, SemanticInternalAttributes, stringifyIO, SubtaskUnwrapError, - resourceCatalog, taskContext, + TaskFromIdentifier, TaskRunContext, TaskRunExecutionResult, TaskRunPromise, - TaskFromIdentifier, - flattenIdempotencyKey, - getEnvVar, } from "@trigger.dev/core/v3"; import { PollOptions, runs } from "./runs.js"; import { tracer } from "./tracer.js"; import type { + AnyOnCatchErrorHookFunction, + AnyOnCleanupHookFunction, + AnyOnCompleteHookFunction, + AnyOnFailureHookFunction, + AnyOnInitHookFunction, + AnyOnMiddlewareHookFunction, + AnyOnResumeHookFunction, + AnyOnStartHookFunction, + AnyOnSuccessHookFunction, + AnyOnWaitHookFunction, AnyRunHandle, AnyRunTypes, AnyTask, + AnyTaskRunResult, BatchByIdAndWaitItem, - BatchByTaskAndWaitItem, BatchByIdItem, + BatchByIdResult, + BatchByTaskAndWaitItem, BatchByTaskItem, BatchByTaskResult, - BatchByIdResult, BatchItem, BatchResult, BatchRunHandle, BatchRunHandleFromTypes, BatchTasksRunHandleFromTypes, BatchTriggerAndWaitItem, + BatchTriggerAndWaitOptions, BatchTriggerOptions, + BatchTriggerTaskV2RequestBody, InferRunTypes, inferSchemaIn, inferToolParameters, @@ -74,9 +88,6 @@ import type { TriggerAndWaitOptions, TriggerApiRequestOptions, TriggerOptions, - AnyTaskRunResult, - BatchTriggerAndWaitOptions, - BatchTriggerTaskV2RequestBody, } from "@trigger.dev/core/v3"; export type { @@ -93,6 +104,7 @@ export type { SerializableJson, Task, TaskBatchOutputHandle, + TaskFromIdentifier, TaskIdentifier, TaskOptions, TaskOutput, @@ -100,7 +112,6 @@ export type { TaskPayload, TaskRunResult, TriggerOptions, - TaskFromIdentifier, }; export { SubtaskUnwrapError, TaskRunPromise }; @@ -184,6 +195,8 @@ export function createTask< }, }; + registerTaskLifecycleHooks(params.id, params); + resourceCatalog.registerTaskMetadata({ id: params.id, description: params.description, @@ -193,13 +206,6 @@ export function createTask< maxDuration: params.maxDuration, fns: { run: params.run, - init: params.init, - cleanup: params.cleanup, - middleware: params.middleware, - handleError: params.handleError, - onSuccess: params.onSuccess, - onFailure: params.onFailure, - onStart: params.onStart, }, }); @@ -316,6 +322,8 @@ export function createSchemaTask< }, }; + registerTaskLifecycleHooks(params.id, params); + resourceCatalog.registerTaskMetadata({ id: params.id, description: params.description, @@ -325,13 +333,6 @@ export function createSchemaTask< maxDuration: params.maxDuration, fns: { run: params.run, - init: params.init, - cleanup: params.cleanup, - middleware: params.middleware, - handleError: params.handleError, - onSuccess: params.onSuccess, - onFailure: params.onFailure, - onStart: params.onStart, parsePayload, }, }); @@ -1559,3 +1560,77 @@ async function handleTaskRunExecutionResult(taskId: TIdentifier, params: TaskOptions) { + if (params.init) { + lifecycleHooks.registerTaskInitHook(taskId, { + fn: params.init as AnyOnInitHookFunction, + }); + } + + if (params.onStart) { + lifecycleHooks.registerTaskStartHook(taskId, { + fn: params.onStart as AnyOnStartHookFunction, + }); + } + + if (params.onFailure) { + lifecycleHooks.registerTaskFailureHook(taskId, { + fn: params.onFailure as AnyOnFailureHookFunction, + }); + } + + if (params.onSuccess) { + lifecycleHooks.registerTaskSuccessHook(taskId, { + fn: params.onSuccess as AnyOnSuccessHookFunction, + }); + } + + if (params.onComplete) { + lifecycleHooks.registerTaskCompleteHook(taskId, { + fn: params.onComplete as AnyOnCompleteHookFunction, + }); + } + + if (params.onWait) { + lifecycleHooks.registerTaskWaitHook(taskId, { + fn: params.onWait as AnyOnWaitHookFunction, + }); + } + + if (params.onResume) { + lifecycleHooks.registerTaskResumeHook(taskId, { + fn: params.onResume as AnyOnResumeHookFunction, + }); + } + + if (params.catchError) { + // We don't need to use an adapter here because catchError is the new version of handleError + lifecycleHooks.registerTaskCatchErrorHook(taskId, { + fn: params.catchError as AnyOnCatchErrorHookFunction, + }); + } + + if (params.handleError) { + lifecycleHooks.registerTaskCatchErrorHook(taskId, { + fn: params.handleError as AnyOnCatchErrorHookFunction, + }); + } + + if (params.middleware) { + lifecycleHooks.registerTaskMiddlewareHook(taskId, { + fn: params.middleware as AnyOnMiddlewareHookFunction, + }); + } + + if (params.cleanup) { + lifecycleHooks.registerTaskCleanupHook(taskId, { + fn: params.cleanup as AnyOnCleanupHookFunction, + }); + } +} diff --git a/packages/trigger-sdk/src/v3/tasks.ts b/packages/trigger-sdk/src/v3/tasks.ts index c38f77187a..a6089d090e 100644 --- a/packages/trigger-sdk/src/v3/tasks.ts +++ b/packages/trigger-sdk/src/v3/tasks.ts @@ -1,3 +1,14 @@ +import { + onStart, + onFailure, + onSuccess, + onComplete, + onWait, + onResume, + onHandleError, + onCatchError, + middleware, +} from "./hooks.js"; import { batchTrigger, batchTriggerAndWait, @@ -46,6 +57,8 @@ export type { TaskFromIdentifier, }; +export type * from "./hooks.js"; + /** Creates a task that can be triggered * @param options - Task options * @example @@ -76,4 +89,14 @@ export const tasks = { batchTrigger, triggerAndWait, batchTriggerAndWait, + onStart, + onFailure, + onSuccess, + onComplete, + onWait, + onResume, + /** @deprecated Use catchError instead */ + handleError: onHandleError, + catchError: onCatchError, + middleware, }; diff --git a/references/hello-world/src/db.ts b/references/hello-world/src/db.ts new file mode 100644 index 0000000000..7c2cd12c15 --- /dev/null +++ b/references/hello-world/src/db.ts @@ -0,0 +1,48 @@ +import { locals } from "@trigger.dev/sdk"; +import { logger, tasks } from "@trigger.dev/sdk"; + +const DbLocal = locals.create<{ connect: () => Promise; disconnect: () => Promise }>( + "db" +); + +export function getDb() { + return locals.getOrThrow(DbLocal); +} + +export function setDb(db: { connect: () => Promise }) { + locals.set(DbLocal, db); +} + +tasks.middleware("db", async ({ ctx, payload, next, task }) => { + const db = locals.set(DbLocal, { + connect: async () => { + logger.info("Connecting to the database"); + }, + disconnect: async () => { + logger.info("Disconnecting from the database"); + }, + }); + + await db.connect(); + + logger.info("Hello, world from BEFORE the next call", { ctx, payload }); + await next(); + + logger.info("Hello, world from AFTER the next call", { ctx, payload }); + + await db.disconnect(); +}); + +tasks.onWait("db", async ({ ctx, payload, task }) => { + logger.info("Hello, world from ON WAIT", { ctx, payload }); + + const db = getDb(); + await db.disconnect(); +}); + +tasks.onResume("db", async ({ ctx, payload, task }) => { + logger.info("Hello, world from ON RESUME", { ctx, payload }); + + const db = getDb(); + await db.connect(); +}); diff --git a/references/hello-world/src/trigger/example.ts b/references/hello-world/src/trigger/example.ts index 81c45ec32c..d1b008f417 100644 --- a/references/hello-world/src/trigger/example.ts +++ b/references/hello-world/src/trigger/example.ts @@ -4,6 +4,8 @@ import { setTimeout } from "timers/promises"; export const helloWorldTask = task({ id: "hello-world", run: async (payload: any, { ctx }) => { + logger.info("Hello, world from the init", { ctx, payload }); + logger.debug("debug: Hello, world!", { payload }); logger.info("info: Hello, world!", { payload }); logger.log("log: Hello, world!", { payload }); @@ -145,3 +147,48 @@ const nonExportedTask = task({ logger.info("Hello, world from the non-exported task", { message: payload.message }); }, }); + +export const hooksTask = task({ + id: "hooks", + run: async (payload: { message: string }, { ctx }) => { + logger.info("Hello, world from the hooks task", { message: payload.message }); + + await wait.for({ seconds: 5 }); + + return { + message: "Hello, world!", + }; + }, + init: async () => { + return { + foobar: "baz", + }; + }, + onWait: async ({ payload, wait, ctx, init }) => { + logger.info("Hello, world from the onWait hook", { payload, init, wait }); + }, + onResume: async ({ payload, wait, ctx, init }) => { + logger.info("Hello, world from the onResume hook", { payload, init, wait }); + }, + onStart: async ({ payload, ctx, init }) => { + logger.info("Hello, world from the onStart hook", { payload, init }); + }, + onSuccess: async ({ payload, output, ctx }) => { + logger.info("Hello, world from the onSuccess hook", { payload, output }); + }, + onFailure: async ({ payload, error, ctx }) => { + logger.info("Hello, world from the onFailure hook", { payload, error }); + }, + onComplete: async ({ ctx, payload, result }) => { + logger.info("Hello, world from the onComplete hook", { payload, result }); + }, + handleError: async ({ payload, error, ctx, retry }) => { + logger.info("Hello, world from the handleError hook", { payload, error, retry }); + }, + catchError: async ({ ctx, payload, error, retry }) => { + logger.info("Hello, world from the catchError hook", { payload, error, retry }); + }, + cleanup: async ({ ctx, payload }) => { + logger.info("Hello, world from the cleanup hook", { payload }); + }, +}); diff --git a/references/hello-world/src/trigger/init.ts b/references/hello-world/src/trigger/init.ts new file mode 100644 index 0000000000..e57bc89838 --- /dev/null +++ b/references/hello-world/src/trigger/init.ts @@ -0,0 +1,41 @@ +import { logger, tasks } from "@trigger.dev/sdk"; +// import { setDb } from "../db.js"; + +tasks.middleware("db", ({ ctx, payload, next }) => { + logger.info("Hello, world from the middleware", { ctx, payload }); + return next(); +}); + +// tasks.onSuccess(({ ctx, payload, output }) => { +// logger.info("Hello, world from the success", { ctx, payload }); +// }); + +// tasks.onComplete(({ ctx, payload, output, error }) => { +// logger.info("Hello, world from the success", { ctx, payload }); +// }); + +// tasks.handleError(({ ctx, payload, error, retry, retryAt, retryDelayInMs }) => { +// logger.info("Hello, world from the success", { ctx, payload }); +// }); + +// tasks.onFailure(({ ctx, payload }) => { +// logger.info("Hello, world from the failure", { ctx, payload }); +// }); + +// tasks.onStart(({ ctx, payload }) => { +// logger.info("Hello, world from the start", { ctx, payload }); + +// setDb({ +// connect: async () => { +// logger.info("Connecting to the database"); +// }, +// }); +// }); + +// tasks.onWait(({ ctx, payload }) => { +// logger.info("Hello, world from the start", { ctx, payload }); +// }); + +// tasks.onResume(({ ctx, payload }) => { +// logger.info("Hello, world from the start", { ctx, payload }); +// });