Improve init lifecycle hook types and fix tabler icons on spans

ericallam · ericallam · commit 17e162310dc5 · 2025-03-25T15:55:42.000Z
diff --git a/apps/webapp/app/components/runs/v3/RunIcon.tsx b/apps/webapp/app/components/runs/v3/RunIcon.tsx
@@ -28,6 +28,8 @@ type SpanNameIcons = {
 const spanNameIcons: SpanNameIcons[] = [{ matcher: /^prisma:/, iconName: "brand-prisma" }];
 
 export function RunIcon({ name, className, spanName }: TaskIconProps) {
+  console.log("spanName", spanName, name);
+
   const spanNameIcon = spanNameIcons.find(({ matcher }) => matcher.test(spanName));
 
   if (spanNameIcon) {
@@ -44,6 +46,9 @@ export function RunIcon({ name, className, spanName }: TaskIconProps) {
   }
 
   if (!name) return <Squares2X2Icon className={cn(className, "text-text-dimmed")} />;
+  if (tablerIcons.has(name)) {
+    return <TablerIcon name={name} className={className} />;
+  }
 
   switch (name) {
     case "task":
diff --git a/packages/core/src/v3/lifecycle-hooks-api.ts b/packages/core/src/v3/lifecycle-hooks-api.ts
@@ -8,6 +8,7 @@ export type {
   OnInitHookFunction,
   AnyOnInitHookFunction,
   RegisteredHookFunction,
+  TaskInitHookParams,
 } from "./lifecycleHooks/types.js";
 
 export * as lifecycleHooksAdapters from "./lifecycleHooks/adapters.js";
diff --git a/packages/core/src/v3/lifecycleHooks/manager.ts b/packages/core/src/v3/lifecycleHooks/manager.ts
@@ -6,20 +6,20 @@ import {
 } from "./types.js";
 
 export class StandardLifecycleHooksManager implements LifecycleHooksManager {
-  private initHooks: Map<string, RegisteredHookFunction<AnyOnInitHookFunction>> = new Map();
-  private taskInitHooks: Map<string, string> = new Map();
+  private globalInitHooks: Map<string, RegisteredHookFunction<AnyOnInitHookFunction>> = new Map();
+  private taskInitHooks: Map<string, RegisteredHookFunction<AnyOnInitHookFunction>> = new Map();
 
   registerGlobalInitHook(hook: RegisterHookFunctionParams<AnyOnInitHookFunction>): 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 ?? hook.fn.name,
+      name: hook.id ?? hook.fn.name ? (hook.fn.name === "" ? undefined : hook.fn.name) : undefined,
       fn: hook.fn,
     };
 
-    this.initHooks.set(id, registeredHook);
+    this.globalInitHooks.set(id, registeredHook);
   }
 
   registerTaskInitHook(
@@ -32,18 +32,15 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager {
       fn: hook.fn,
     };
 
-    this.initHooks.set(registeredHook.id, registeredHook);
-    this.taskInitHooks.set(taskId, registeredHook.id);
+    this.taskInitHooks.set(taskId, registeredHook);
   }
 
   getTaskInitHook(taskId: string): AnyOnInitHookFunction | undefined {
-    const hookId = this.taskInitHooks.get(taskId);
-    if (!hookId) return undefined;
-    return this.initHooks.get(hookId)?.fn;
+    return this.taskInitHooks.get(taskId)?.fn;
   }
 
   getGlobalInitHooks(): RegisteredHookFunction<AnyOnInitHookFunction>[] {
-    return Array.from(this.initHooks.values());
+    return Array.from(this.globalInitHooks.values());
   }
 }
 
@@ -69,5 +66,9 @@ export class NoopLifecycleHooksManager implements LifecycleHooksManager {
 }
 
 function generateHookId(hook: RegisterHookFunctionParams<any>): string {
-  return hook.id ?? hook.fn.name ?? hook.fn.toString();
+  return hook.id ?? hook.fn.name
+    ? hook.fn.name === ""
+      ? hook.fn.toString()
+      : hook.fn.name
+    : hook.fn.toString();
 }
diff --git a/packages/core/src/v3/lifecycleHooks/types.ts b/packages/core/src/v3/lifecycleHooks/types.ts
@@ -1,13 +1,17 @@
 import { TaskRunContext } from "../schemas/index.js";
 
-export type OnInitHookFunction<TPayload, TInitOutput> = (params: {
+export type TaskInitHookParams<TPayload = unknown> = {
   ctx: TaskRunContext;
   payload: TPayload;
   task: string;
   signal?: AbortSignal;
-}) => TInitOutput | undefined | void | Promise<TInitOutput | undefined | void>;
+};
+
+export type OnInitHookFunction<TPayload, TInitOutput extends Record<string, unknown>> = (
+  params: TaskInitHookParams<TPayload>
+) => TInitOutput | undefined | void | Promise<TInitOutput | undefined | void>;
 
-export type AnyOnInitHookFunction = OnInitHookFunction<unknown, unknown>;
+export type AnyOnInitHookFunction = OnInitHookFunction<unknown, Record<string, unknown>>;
 
 export type RegisterHookFunctionParams<THookFunction extends (params: any) => any> = {
   id?: string;
diff --git 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/workers/taskExecutor.ts b/packages/core/src/v3/workers/taskExecutor.ts
@@ -3,7 +3,13 @@ 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 { lifecycleHooks, runMetadata, TriggerConfig, waitUntil } from "../index.js";
+import {
+  flattenAttributes,
+  lifecycleHooks,
+  runMetadata,
+  TriggerConfig,
+  waitUntil,
+} from "../index.js";
 import { recordSpanException, TracingSDK } from "../otel/index.js";
 import { runTimelineMetrics } from "../run-timeline-metrics-api.js";
 import {
@@ -288,58 +294,87 @@ export class TaskExecutor {
     return this._tracer.startActiveSpan(
       "hooks.init",
       async (span) => {
-        return await runTimelineMetrics.measureMetric("trigger.dev/execution", "init", async () => {
-          // Store global hook results in an array
-          const globalResults = [];
-          for (const hook of globalInitHooks) {
-            const result = await this._tracer.startActiveSpan(
-              hook.name ?? "global",
-              async (span) => {
-                return await hook.fn({ payload, ctx, signal, task: this.task.id });
-              },
-              {
-                attributes: {
-                  [SemanticInternalAttributes.STYLE_ICON]: "function",
+        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 result = await this._tracer.startActiveSpan(
+                hook.name ?? "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]: "tabler-function",
+                  },
+                }
+              );
+              // Only include object results
+              if (result && typeof result === "object" && !Array.isArray(result)) {
+                globalResults.push(result);
               }
-            );
-            // Only include object results
-            if (result && typeof result === "object" && !Array.isArray(result)) {
-              globalResults.push(result);
             }
-          }
 
-          // Merge all global results into a single object
-          const mergedGlobalResults = Object.assign({}, ...globalResults);
+            // Merge all global results into a single object
+            const mergedGlobalResults = Object.assign({}, ...globalResults);
 
-          if (taskInitHook) {
-            const taskResult = await this._tracer.startActiveSpan(
-              "task",
-              async (span) => {
-                return await taskInitHook({ payload, ctx, signal, task: this.task.id });
-              },
-              {
-                attributes: {
-                  [SemanticInternalAttributes.STYLE_ICON]: "function",
+            if (taskInitHook) {
+              const taskResult = await this._tracer.startActiveSpan(
+                "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]: "tabler-function",
+                  },
+                }
+              );
+
+              // Only merge if taskResult is an object
+              if (taskResult && typeof taskResult === "object" && !Array.isArray(taskResult)) {
+                return { ...mergedGlobalResults, ...taskResult };
               }
-            );
 
-            // 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;
             }
 
-            // If taskResult isn't an object, return global results
             return mergedGlobalResults;
           }
+        );
 
-          return mergedGlobalResults;
-        });
+        if (result && typeof result === "object" && !Array.isArray(result)) {
+          span.setAttributes(flattenAttributes(result));
+
+          return result;
+        }
+
+        return;
       },
       {
         attributes: {
-          [SemanticInternalAttributes.STYLE_ICON]: "function",
+          [SemanticInternalAttributes.STYLE_ICON]: "tabler-function",
+          [SemanticInternalAttributes.COLLAPSED]: true,
         },
       }
     );
diff --git a/packages/trigger-sdk/src/v3/hooks.ts b/packages/trigger-sdk/src/v3/hooks.ts
@@ -1,4 +1,11 @@
-import { lifecycleHooks, type AnyOnInitHookFunction } from "@trigger.dev/core/v3";
+import {
+  lifecycleHooks,
+  type AnyOnInitHookFunction,
+  type TaskInitHookParams,
+  type OnInitHookFunction,
+} from "@trigger.dev/core/v3";
+
+export type { AnyOnInitHookFunction, TaskInitHookParams, OnInitHookFunction };
 
 export function onInit(name: string, fn: AnyOnInitHookFunction): void;
 export function onInit(fn: AnyOnInitHookFunction): void;
diff --git a/packages/trigger-sdk/src/v3/tasks.ts b/packages/trigger-sdk/src/v3/tasks.ts
@@ -47,6 +47,8 @@ export type {
   TaskFromIdentifier,
 };
 
+export type * from "./hooks.js";
+
 /** Creates a task that can be triggered
  * @param options - Task options
  * @example 
diff --git a/references/hello-world/src/trigger/example.ts b/references/hello-world/src/trigger/example.ts
@@ -3,8 +3,15 @@ import { setTimeout } from "timers/promises";
 
 export const helloWorldTask = task({
   id: "hello-world",
-  run: async (payload: any, { ctx }) => {
-    logger.debug("debug: Hello, world!", { payload });
+  init: async (payload, { ctx }) => {
+    logger.info("Hello, world from the init", { ctx, payload });
+
+    return {
+      foobar: "baz",
+    };
+  },
+  run: async (payload: any, { ctx, init }) => {
+    logger.debug("debug: Hello, world!", { payload, init });
     logger.info("info: Hello, world!", { payload });
     logger.log("log: Hello, world!", { payload });
     logger.warn("warn: Hello, world!", { payload });
diff --git a/references/hello-world/src/trigger/init.ts b/references/hello-world/src/trigger/init.ts
@@ -1,4 +1,4 @@
-import { logger, tasks } from "@trigger.dev/sdk";
+import { logger, type TaskInitHookParams, tasks } from "@trigger.dev/sdk";
 
 // tasks.onSuccess(({ ctx, payload, output }) => {
 //   logger.info("Hello, world from the success", { ctx, payload });
@@ -28,6 +28,18 @@ import { logger, tasks } from "@trigger.dev/sdk";
 //   logger.info("Hello, world from the start", { ctx, payload });
 // });
 
-tasks.onInit("logging", ({ ctx, payload, task }) => {
+tasks.onInit(({ ctx, payload, task }) => {
   logger.info("Hello, world from the init", { ctx, payload, task });
+
+  return {
+    foo: "bar",
+  };
+});
+
+tasks.onInit(({ ctx, payload, task }) => {
+  logger.info("Hello, world from the init 2", { ctx, payload, task });
+
+  return {
+    bar: "baz",
+  };
 });
