feat(sdk): add onBoot lifecycle hook to chat.agent

Author: ericallam

.changeset/chat-agent-on-boot-hook.md

@@ -0,0 +1,21 @@
+---
+"@trigger.dev/sdk": minor
+---
+
+Adds `onBoot` to `chat.agent` — a lifecycle hook that fires once per worker process picking up the chat. Runs for the initial run, preloaded runs, AND reactive continuation runs (post-cancel, crash, `endRun`, `requestUpgrade`, OOM retry), before any other hook. Use it to initialize `chat.local`, open per-process resources, or re-hydrate state from your DB on continuation — anywhere the SAME run picking up after suspend/resume isn't enough.
+
+```ts
+const userContext = chat.local<{ name: string; plan: string }>({ id: "userContext" });
+
+export const myChat = chat.agent({
+  id: "my-chat",
+  onBoot: async ({ clientData, continuation }) => {
+    const user = await db.user.findUnique({ where: { id: clientData.userId } });
+    userContext.init({ name: user.name, plan: user.plan });
+  },
+  run: async ({ messages, signal }) =>
+    streamText({ model: openai("gpt-4o"), messages, abortSignal: signal }),
+});
+```
+
+If you previously initialized `chat.local` in `onChatStart`, move it to `onBoot` — `onChatStart` is once-per-chat and won't fire on a continuation, leaving `chat.local` uninitialized when `run()` tries to use it. See the upgrade guide for the migration pattern.

packages/trigger-sdk/src/v3/ai.ts

@@ -3583,6 +3583,50 @@ export type PreloadEvent<TClientData = unknown> = {
   writer: ChatWriter;
 };
 
+/**
+ * Event passed to the `onBoot` callback.
+ *
+ * Fires once at the start of every run boot — for the initial first-message
+ * run, for preloaded runs, AND for reactive continuation runs (post-cancel,
+ * post-crash, post-`endRun`, `chat.requestUpgrade`, OOM-retry attempts).
+ * Does NOT fire when the SAME run resumes from snapshot via the
+ * idle-window suspend/resume path — use `onChatResume` for that.
+ *
+ * Use this for per-process setup that needs to run every time a fresh
+ * worker picks up the chat: initialize `chat.local` state, open
+ * per-process resources (DB connections, sandboxes, etc.), or
+ * re-hydrate customer state from your DB on continuation.
+ *
+ * Ordering:
+ *   - First message of a chat: `onBoot` → `onChatStart` → `onTurnStart` → `run()`
+ *   - Preloaded run:           `onBoot` → `onPreload` → (wait) → `onChatStart` → ...
+ *   - Continuation run:        `onBoot` → (wait for first message) → `onTurnStart` → `run()`
+ *                              (`onChatStart` does NOT fire on continuation runs)
+ */
+export type BootEvent<TClientData = unknown> = {
+  /** Task run context — same as `task({ run })` second-argument `ctx`. */
+  ctx: TaskRunContext;
+  /** The unique identifier for the chat session. */
+  chatId: string;
+  /** The Trigger.dev run ID for this run boot. */
+  runId: string;
+  /** A scoped access token for this chat run. Persist this for frontend reconnection. */
+  chatAccessToken: string;
+  /** Custom data from the frontend (passed via `metadata` on `sendMessage()` or the transport). */
+  clientData: TClientData;
+  /**
+   * True when this run is a reactive continuation — the prior run died
+   * (cancel / crash / `endRun` / `requestUpgrade` / OOM retry) and this
+   * fresh worker is taking over the chat. Branch on this to re-load
+   * customer-owned state from your DB.
+   */
+  continuation: boolean;
+  /** Public id of the prior run when `continuation` is true. */
+  previousRunId?: string;
+  /** Whether this run was triggered as a preload. */
+  preloaded: boolean;
+};
+
 /**
  * Event passed to the `onChatStart` callback.
  *
@@ -4070,6 +4114,36 @@ export type ChatAgentOptions<
    */
   run: (payload: ChatTaskRunPayload<inferSchemaOut<TClientDataSchema>>) => Promise<unknown>;
 
+  /**
+   * Called once at the start of every run boot — for the initial run, for
+   * preloaded runs, AND for reactive continuation runs (post-cancel /
+   * crash / `endRun` / `requestUpgrade` / OOM retry).
+   *
+   * Use this for per-process setup that needs to run every time a fresh
+   * worker picks up the chat: initialize `chat.local` state, open
+   * per-process resources, or re-hydrate customer state from your DB
+   * on continuation. Branch on `continuation` to decide whether to load
+   * existing state vs. start fresh.
+   *
+   * Fires BEFORE `onPreload` (preloaded path) and `onChatStart`
+   * (first-message path), so `chat.local` is safe to read in those hooks.
+   *
+   * Does NOT fire when the same run resumes from snapshot via the
+   * idle-window suspend/resume path — use `onChatResume` for that.
+   *
+   * @example
+   * ```ts
+   * onBoot: async ({ chatId, clientData, continuation }) => {
+   *   const user = await db.user.findFirst({ where: { id: clientData.userId } });
+   *   userContext.init({ name: user.name, plan: user.plan });
+   *   if (continuation) {
+   *     // re-hydrate any per-chat in-memory state from your DB
+   *   }
+   * }
+   * ```
+   */
+  onBoot?: (event: BootEvent<inferSchemaOut<TClientDataSchema>>) => Promise<void> | void;
+
   /**
    * Called when a preloaded run starts, before the first message arrives.
    *
@@ -4618,6 +4692,7 @@ function chatAgent<
   const {
     run: userRun,
     clientDataSchema,
+    onBoot,
     onPreload,
     onChatStart,
     onValidateMessages,
@@ -4887,6 +4962,74 @@ function chatAgent<
       let stopSub: { off: () => void } | undefined;
 
       try {
+        // ── onBoot — fires once per fresh worker process picking up the chat.
+        //
+        // Runs BEFORE `onPreload`, `onChatStart`, the continuation-wait
+        // branch, and any turn. This is the hook customers use to do
+        // per-process setup that has to repeat for every boot (initial
+        // run, preloaded run, AND reactive continuation): `chat.local`
+        // initialization, opening per-process resources, re-hydrating
+        // customer state from their DB on continuation.
+        //
+        // `onChatStart` is once-per-chat (`!couldHavePriorState` gate);
+        // when the prior run dies and a fresh worker boots, `onChatStart`
+        // does NOT fire — which is why customers initializing
+        // `chat.local` only in `onChatStart` previously crashed in
+        // `run()` on continuation. `onBoot` fills that gap.
+        if (onBoot) {
+          const bootClientData = (
+            parseClientData ? await parseClientData(payload.metadata) : payload.metadata
+          ) as inferSchemaOut<TClientDataSchema>;
+
+          let bootAccessToken = "";
+          const bootRunId = ctx.run.id;
+          if (bootRunId) {
+            try {
+              bootAccessToken = await auth.createPublicToken({
+                scopes: {
+                  read: {
+                    runs: bootRunId,
+                    sessions: payload.chatId,
+                  },
+                  write: {
+                    inputStreams: bootRunId,
+                    sessions: payload.chatId,
+                  },
+                },
+                expirationTime: chatAccessTokenTTL,
+              });
+            } catch {
+              // Token mint is best-effort — customers can re-mint from
+              // their own backend if they need one.
+            }
+          }
+
+          await tracer.startActiveSpan(
+            "onBoot()",
+            async () => {
+              await onBoot({
+                ctx,
+                chatId: payload.chatId,
+                runId: bootRunId,
+                chatAccessToken: bootAccessToken,
+                clientData: bootClientData,
+                continuation,
+                previousRunId,
+                preloaded,
+              });
+            },
+            {
+              attributes: {
+                [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart",
+                [SemanticInternalAttributes.COLLAPSED]: true,
+                "chat.id": payload.chatId,
+                "chat.continuation": continuation,
+                "chat.preloaded": preloaded,
+              },
+            }
+          );
+        }
+
         // Handle preloaded runs — fire onPreload, then wait for the first real message
         if (preloaded) {
           if (activeSpan) {
@@ -6810,6 +6953,11 @@ export interface ChatBuilder<
     schema: TSchema;
   }): ChatBuilder<TUIMessage, TSchema>;
 
+  /** Register a builder-level `onBoot` hook. Runs before the task-level hook if both are set. */
+  onBoot(
+    fn: (event: BootEvent<inferSchemaOut<TClientDataSchema>>) => Promise<void> | void
+  ): ChatBuilder<TUIMessage, TClientDataSchema>;
+
   /** Register a builder-level `onPreload` hook. Runs before the task-level hook if both are set. */
   onPreload(
     fn: (event: PreloadEvent<inferSchemaOut<TClientDataSchema>>) => Promise<void> | void
@@ -6894,6 +7042,7 @@ export interface ChatBuilder<
 
 /** @internal */
 type ChatBuilderHooks = {
+  onBoot?: (event: any) => Promise<void> | void;
   onPreload?: (event: any) => Promise<void> | void;
   onChatStart?: (event: any) => Promise<void> | void;
   onTurnStart?: (event: any) => Promise<void> | void;
@@ -6942,6 +7091,14 @@ function createChatBuilder<
       });
     },
 
+    onBoot(
+      fn: (event: BootEvent<inferSchemaOut<TClientDataSchema>>) => Promise<void> | void
+    ) {
+      return createChatBuilder<TUIMessage, TClientDataSchema>({
+        ...config,
+        hooks: { ...config.hooks, onBoot: fn },
+      });
+    },
     onPreload(
       fn: (event: PreloadEvent<inferSchemaOut<TClientDataSchema>>) => Promise<void> | void
     ) {
@@ -7025,6 +7182,7 @@ function createChatBuilder<
         ...options,
         ...(config.clientDataSchema ? { clientDataSchema: config.clientDataSchema } : {}),
         uiMessageStreamOptions: mergedUiStream,
+        onBoot: composeHooks(config.hooks.onBoot, options.onBoot),
         onPreload: composeHooks(config.hooks.onPreload, options.onPreload),
         onChatStart: composeHooks(config.hooks.onChatStart, options.onChatStart),
         onTurnStart: composeHooks(config.hooks.onTurnStart, options.onTurnStart),
@@ -8307,8 +8465,11 @@ export type ChatLocal<T extends Record<string, unknown>> = T & {
 /**
  * Creates a per-run typed data object accessible from anywhere during task execution.
  *
- * Declare at module level, then initialize inside a lifecycle hook (e.g. `onChatStart`)
- * using `chat.initLocal()`. Properties are accessible directly via the Proxy.
+ * Declare at module level, then initialize inside `onBoot` (recommended — fires
+ * on every fresh worker including continuation runs). Do NOT initialize in
+ * `onChatStart` alone: `onChatStart` only fires on the chat's very first
+ * message, so `chat.local` would be uninitialized on continuation runs and
+ * `run()` would throw.
  *
  * Multiple locals can coexist — each gets its own isolated run-scoped storage.
  *
@@ -8325,7 +8486,7 @@ export type ChatLocal<T extends Record<string, unknown>> = T & {
  *
  * export const myChat = chat.agent({
  *   id: "my-chat",
- *   onChatStart: async ({ clientData }) => {
+ *   onBoot: async ({ clientData }) => {
  *     const prefs = await db.prefs.findUnique({ where: { userId: clientData.userId } });
  *     userPrefs.init(prefs ?? { theme: "dark", language: "en" });
  *     gameState.init({ score: 0, streak: 0 });
@@ -8384,7 +8545,9 @@ function chatLocal<T extends Record<string, unknown>>(options: { id: string }):
             current = locals.get(localKey);
           }
           if (current === undefined) {
-            throw new Error("local.get() called before initialization. Call local.init() first.");
+            throw new Error(
+              "local.get() called before initialization. Call local.init() in onBoot (recommended — fires on every fresh worker including continuation runs) or run() first."
+            );
           }
           return { ...current };
         };
@@ -8419,7 +8582,8 @@ function chatLocal<T extends Record<string, unknown>>(options: { id: string }):
       if (current === undefined) {
         throw new Error(
           "chat.local can only be modified after initialization. " +
-          "Call local.init() in onChatStart or run() first."
+          "Call local.init() in onBoot (recommended — fires on every fresh worker including continuation runs) or run() first. " +
+          "If you previously initialized in onChatStart, move it to onBoot — onChatStart only fires on the chat's very first message and will not run on a continuation."
         );
       }
       locals.set(localKey, { ...current, [prop]: value });