Support async onTriggeredRun callbacks in chat transport

Co-authored-by: Eric Allam <eric@trigger.dev>

Author: cursoragent

packages/ai/README.md

@@ -116,6 +116,9 @@ class MemoryStore implements TriggerChatRunStore {
 }
 ```
 
+`onTriggeredRun` can also be async, which is useful for persisting run IDs before
+the chat stream is consumed.
+
 ## `ai.tool(...)` example
 
 ```ts

packages/ai/src/chatTransport.test.ts

@@ -628,6 +628,7 @@ describe("TriggerChatTransport", function () {
 
   it("supports creating transport with factory function", async function () {
     let observedRunId: string | undefined;
+    let callbackCompleted = false;
 
     const server = await startServer(function (req, res) {
       if (req.method === "POST" && req.url === "/api/v1/tasks/chat-task/trigger") {
@@ -666,8 +667,10 @@ describe("TriggerChatTransport", function () {
       stream: "chat-stream",
       accessToken: "pk_trigger",
       baseURL: server.url,
-      onTriggeredRun: function onTriggeredRun(state) {
+      onTriggeredRun: async function onTriggeredRun(state) {
+        await sleep(1);
         observedRunId = state.runId;
+        callbackCompleted = true;
       },
     });
 
@@ -682,6 +685,7 @@ describe("TriggerChatTransport", function () {
     const chunks = await readChunks(stream);
     expect(chunks).toHaveLength(2);
     expect(observedRunId).toBe("run_factory");
+    expect(callbackCompleted).toBe(true);
   });
 
   it("cleans run store state when stream completes", async function () {

packages/ai/src/chatTransport.ts

@@ -15,6 +15,7 @@ import type {
   UIMessageChunk,
 } from "ai";
 import type {
+  TriggerChatOnTriggeredRun,
   TriggerChatPayloadMapper,
   TriggerChatRunState,
   TriggerChatRunStore,
@@ -75,7 +76,7 @@ type TriggerChatTransportCommonOptions<
     | TriggerOptions
     | TriggerChatTriggerOptionsResolver<UI_MESSAGE>;
   runStore?: TriggerChatRunStore;
-  onTriggeredRun?: (state: TriggerChatRunState) => void;
+  onTriggeredRun?: TriggerChatOnTriggeredRun;
 };
 
 type TriggerChatTransportMapperRequirement<
@@ -129,7 +130,7 @@ export class TriggerChatTransport<
   private readonly baseURL: string;
   private readonly previewBranch: string | undefined;
   private readonly requestOptions: ApiRequestOptions | undefined;
-  private readonly onTriggeredRun: ((state: TriggerChatRunState) => void) | undefined;
+  private readonly onTriggeredRun: TriggerChatOnTriggeredRun | undefined;
 
   constructor(options: TriggerChatTransportOptions<UI_MESSAGE, PAYLOAD>) {
     this.task = options.task;
@@ -179,7 +180,7 @@ export class TriggerChatTransport<
     await this.runStore.set(runState);
 
     if (this.onTriggeredRun) {
-      this.onTriggeredRun(runState);
+      await this.onTriggeredRun(runState);
     }
 
     const stream = await this.fetchRunStream(runState, options.abortSignal);

packages/ai/src/chatTransport.types.test.ts

@@ -81,6 +81,9 @@ it("accepts async payload mappers and trigger option resolvers", function () {
         queue: `queue-${request.chatId}`,
       };
     },
+    onTriggeredRun: async function onTriggeredRun(_state) {
+      return;
+    },
   };
 
   expectTypeOf(options).toBeObject();

packages/ai/src/index.ts

@@ -7,6 +7,7 @@ export {
 } from "./chatTransport.js";
 export type {
   TriggerChatPayloadMapper,
+  TriggerChatOnTriggeredRun,
   TriggerChatRunState,
   TriggerChatRunStore,
   TriggerChatStream,

packages/ai/src/types.ts

@@ -61,6 +61,10 @@ export type TriggerChatTriggerOptionsResolver<
   request: TriggerChatTransportRequest<UI_MESSAGE>
 ) => MaybePromise<TriggerOptions | undefined>;
 
+export type TriggerChatOnTriggeredRun = (
+  state: TriggerChatRunState
+) => MaybePromise<void>;
+
 export type TriggerChatStream<
   UI_MESSAGE extends UIMessage = UIMessage,
 > =