Get it all working

Author: ericallam

packages/cli-v3/src/entryPoints/managed-run-worker.ts

@@ -18,6 +18,8 @@ import {
   waitUntil,
   apiClientManager,
   runTimelineMetrics,
+  lifecycleHooks,
+  lifecycleHooksAdapters,
   localsAPI,
 } from "@trigger.dev/core/v3";
 import { TriggerTracer } from "@trigger.dev/core/v3/tracer";
@@ -40,6 +42,7 @@ import {
   StandardWaitUntilManager,
   ManagedRuntimeManager,
   StandardRunTimelineMetricsManager,
+  StandardLifecycleHooksManager,
   StandardLocalsManager,
 } from "@trigger.dev/core/v3/workers";
 import { ZodIpcConnection } from "@trigger.dev/core/v3/zodIpc";
@@ -98,6 +101,9 @@ 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();
@@ -185,12 +191,46 @@ async function bootstrap() {
 
   logger.setGlobalTaskLogger(otelTaskLogger);
 
+  if (config.init) {
+    lifecycleHooks.registerGlobalInitHook({
+      id: "config",
+      fn: lifecycleHooksAdapters.createInitHookAdapter(config.init),
+    });
+  }
+
+  if (config.onStart) {
+    lifecycleHooks.registerGlobalStartHook({
+      id: "config",
+      fn: lifecycleHooksAdapters.createStartHookAdapter(config.onStart),
+    });
+  }
+
+  if (config.onSuccess) {
+    lifecycleHooks.registerGlobalSuccessHook({
+      id: "config",
+      fn: lifecycleHooksAdapters.createSuccessHookAdapter(config.onSuccess),
+    });
+  }
+
+  if (config.onFailure) {
+    lifecycleHooks.registerGlobalFailureHook({
+      id: "config",
+      fn: lifecycleHooksAdapters.createFailureHookAdapter(config.onFailure),
+    });
+  }
+
+  if (handleError) {
+    lifecycleHooks.registerGlobalCatchErrorHook({
+      id: "config",
+      fn: lifecycleHooksAdapters.createHandleErrorHookAdapter(handleError),
+    });
+  }
+
   return {
     tracer,
     tracingSDK,
     consoleInterceptor,
     config,
-    handleErrorFn: handleError,
     workerManifest,
   };
 }
@@ -232,7 +272,7 @@ const zodIpc = new ZodIpcConnection({
       }
 
       try {
-        const { tracer, tracingSDK, consoleInterceptor, config, handleErrorFn, workerManifest } =
+        const { tracer, tracingSDK, consoleInterceptor, config, workerManifest } =
           await bootstrap();
 
         _tracingSDK = tracingSDK;
@@ -336,8 +376,7 @@ const zodIpc = new ZodIpcConnection({
           tracer,
           tracingSDK,
           consoleInterceptor,
-          config,
-          handleErrorFn,
+          retries: config.retries,
         });
 
         try {
@@ -355,42 +394,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);
 

packages/core/src/v3/locals-api.ts

@@ -19,8 +19,9 @@ export const locals = {
     }
     return value;
   },
-  set<T>(key: LocalsKey<T>, value: T): void {
+  set<T>(key: LocalsKey<T>, value: T): T {
     localsAPI.setLocal(key, value);
+    return value;
   },
 };
 

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(

references/hello-world/src/db.ts

@@ -1,7 +1,9 @@
 import { locals } from "@trigger.dev/sdk";
 import { logger, tasks } from "@trigger.dev/sdk";
 
-const DbLocal = locals.create<{ connect: () => Promise<void> }>("db");
+const DbLocal = locals.create<{ connect: () => Promise<void>; disconnect: () => Promise<void> }>(
+  "db"
+);
 
 export function getDb() {
   return locals.getOrThrow(DbLocal);
@@ -11,13 +13,22 @@ export function setDb(db: { connect: () => Promise<void> }) {
   locals.set(DbLocal, db);
 }
 
-// tasks.middleware("db", ({ ctx, payload, next, task }) => {
-//   locals.set(DbLocal, {
-//     connect: async () => {
-//       logger.info("Connecting to the database");
-//     },
-//   });
+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");
+    },
+  });
 
-//   logger.info("Hello, world from the middleware", { ctx, payload });
-//   return next();
-// });
+  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();
+});