Support async payload and trigger option resolvers

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

Author: cursoragent

packages/ai/README.md

@@ -61,7 +61,7 @@ type Payload = TriggerChatTransportPayload<UIMessage>;
 
 ## Custom payload mapping
 
-If your task expects a custom payload shape, provide `payloadMapper`:
+If your task expects a custom payload shape, provide `payloadMapper` (sync or async):
 
 ```ts
 import { TriggerChatTransport } from "@trigger.dev/ai";
@@ -73,7 +73,9 @@ const transport = new TriggerChatTransport<
 >({
   task: "ai-chat-custom",
   accessToken: "pk_...",
-  payloadMapper: function payloadMapper(request) {
+  payloadMapper: async function payloadMapper(request) {
+    await Promise.resolve();
+
     const firstPart = request.messages[0]?.parts[0];
 
     return {
@@ -87,6 +89,9 @@ const transport = new TriggerChatTransport<
 });
 ```
 
+`triggerOptions` can also be a function (sync or async), which gives you access to
+`chatId`, messages, and request context to compute queueing/idempotency options.
+
 ## Optional persistent run state
 
 `TriggerChatTransport` supports custom run stores (including async implementations) to persist reconnect state:

packages/ai/src/chatTransport.test.ts

@@ -323,7 +323,9 @@ describe("TriggerChatTransport", function () {
       stream: "chat-stream",
       accessToken: "pk_trigger",
       baseURL: server.url,
-      payloadMapper: function payloadMapper(request) {
+      payloadMapper: async function payloadMapper(request) {
+        await sleep(1);
+
         const firstMessage = request.messages[0];
         const firstPart = firstMessage?.parts[0];
         const prompt =
@@ -337,7 +339,9 @@ describe("TriggerChatTransport", function () {
           sourceHeader: request.request.headers?.["x-source"],
         };
       },
-      triggerOptions: function triggerOptions(request) {
+      triggerOptions: async function triggerOptions(request) {
+        await sleep(1);
+
         receivedResolverChatId = request.chatId;
         receivedResolverHeader = request.request.headers?.["x-source"];
 

packages/ai/src/chatTransport.ts

@@ -160,7 +160,7 @@ export class TriggerChatTransport<
     } & ChatRequestOptions
   ): Promise<ReadableStream<UIMessageChunk>> {
     const transportRequest = createTransportRequest<UI_MESSAGE>(options);
-    const payload = this.payloadMapper(transportRequest);
+    const payload = await this.payloadMapper(transportRequest);
     const triggerOptions = await resolveTriggerOptions<UI_MESSAGE>(
       this.triggerOptions,
       transportRequest
@@ -430,7 +430,7 @@ async function resolveTriggerOptions<UI_MESSAGE extends UIMessage>(
   }
 
   if (typeof options === "function") {
-    return options(request);
+    return await options(request);
   }
 
   return options;

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

@@ -63,6 +63,29 @@ it("types mapper input with rich request context", function () {
   expectTypeOf(options.payloadMapper).toBeFunction();
 });
 
+it("accepts async payload mappers and trigger option resolvers", function () {
+  const options: TriggerChatTransportOptions<
+    UIMessage,
+    { prompt: string; chatId: string }
+  > = {
+    task: "ai-chat",
+    accessToken: "pk_test",
+    payloadMapper: async function payloadMapper(request) {
+      return {
+        prompt: request.chatId,
+        chatId: request.chatId,
+      };
+    },
+    triggerOptions: async function triggerOptions(request) {
+      return {
+        queue: `queue-${request.chatId}`,
+      };
+    },
+  };
+
+  expectTypeOf(options).toBeObject();
+});
+
 it("accepts typed stream definition objects", function () {
   const typedStream = {
     id: "chat-stream",

packages/ai/src/types.ts

@@ -48,16 +48,18 @@ export type TriggerChatTaskContext<
   streamKey: string;
 };
 
+type MaybePromise<T> = T | Promise<T>;
+
 export type TriggerChatPayloadMapper<
   UI_MESSAGE extends UIMessage = UIMessage,
   PAYLOAD = TriggerChatTransportPayload<UI_MESSAGE>,
-> = (request: TriggerChatTransportRequest<UI_MESSAGE>) => PAYLOAD;
+> = (request: TriggerChatTransportRequest<UI_MESSAGE>) => MaybePromise<PAYLOAD>;
 
 export type TriggerChatTriggerOptionsResolver<
   UI_MESSAGE extends UIMessage = UIMessage,
 > = (
   request: TriggerChatTransportRequest<UI_MESSAGE>
-) => TriggerOptions | undefined;
+) => MaybePromise<TriggerOptions | undefined>;
 
 export type TriggerChatStream<
   UI_MESSAGE extends UIMessage = UIMessage,