From 05232a2e673fcedcc29debd5f0811ef228685387 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 15 Feb 2026 12:54:20 +0000 Subject: [PATCH] feat: add chat transport and AI chat helpers to @trigger.dev/sdk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two new subpath exports: @trigger.dev/sdk/chat (frontend, browser-safe): - TriggerChatTransport — ChatTransport implementation for useChat - createChatTransport() — factory function - TriggerChatTransportOptions type @trigger.dev/sdk/ai (backend, adds to existing ai.tool/ai.currentToolOptions): - chatTask() — pre-typed task wrapper with auto-pipe - pipeChat() — pipe StreamTextResult to realtime stream - CHAT_STREAM_KEY constant - ChatTaskPayload type - ChatTaskOptions type - PipeChatOptions type Co-authored-by: Eric Allam --- packages/ai/src/chatTask.ts | 132 +++++++++++++ packages/ai/src/pipeChat.ts | 137 +++++++++++++ packages/ai/src/types.ts | 22 +-- packages/trigger-sdk/package.json | 17 +- packages/trigger-sdk/src/v3/ai.ts | 242 +++++++++++++++++++++++ packages/trigger-sdk/src/v3/chat.ts | 294 ++++++++++++++++++++++++++++ 6 files changed, 832 insertions(+), 12 deletions(-) create mode 100644 packages/ai/src/chatTask.ts create mode 100644 packages/ai/src/pipeChat.ts create mode 100644 packages/trigger-sdk/src/v3/chat.ts diff --git a/packages/ai/src/chatTask.ts b/packages/ai/src/chatTask.ts new file mode 100644 index 000000000..7f3eb9261 --- /dev/null +++ b/packages/ai/src/chatTask.ts @@ -0,0 +1,132 @@ +import { task as createTask } from "@trigger.dev/sdk"; +import type { Task } from "@trigger.dev/core/v3"; +import type { ChatTaskPayload } from "./types.js"; +import { pipeChat } from "./pipeChat.js"; + +/** + * Options for defining a chat task. + * + * This is a simplified version of the standard task options with the payload + * pre-typed as `ChatTaskPayload`. + */ +export type ChatTaskOptions = { + /** Unique identifier for the task */ + id: TIdentifier; + + /** Optional description of the task */ + description?: string; + + /** Retry configuration */ + retry?: { + maxAttempts?: number; + factor?: number; + minTimeoutInMs?: number; + maxTimeoutInMs?: number; + randomize?: boolean; + }; + + /** Queue configuration */ + queue?: { + name?: string; + concurrencyLimit?: number; + }; + + /** Machine preset for the task */ + machine?: { + preset?: string; + }; + + /** Maximum duration in seconds */ + maxDuration?: number; + + /** + * The main run function for the chat task. + * + * Receives a `ChatTaskPayload` with the conversation messages, chat session ID, + * and trigger type. + * + * **Auto-piping:** If this function returns a value that has a `.toUIMessageStream()` method + * (like a `StreamTextResult` from `streamText()`), the stream will automatically be piped + * to the frontend via the chat realtime stream. If you need to pipe from deeper in your + * code, use `pipeChat()` instead and don't return the result. + */ + run: (payload: ChatTaskPayload) => Promise; +}; + +/** + * An object that has a `toUIMessageStream()` method, like the result of `streamText()`. + */ +type UIMessageStreamable = { + toUIMessageStream: (...args: any[]) => AsyncIterable | ReadableStream; +}; + +function isUIMessageStreamable(value: unknown): value is UIMessageStreamable { + return ( + typeof value === "object" && + value !== null && + "toUIMessageStream" in value && + typeof (value as any).toUIMessageStream === "function" + ); +} + +/** + * Creates a Trigger.dev task pre-configured for AI SDK chat. + * + * This is a convenience wrapper around `task()` from `@trigger.dev/sdk` that: + * - **Pre-types the payload** as `ChatTaskPayload` — no manual typing needed + * - **Auto-pipes the stream** if the `run` function returns a `StreamTextResult` + * + * Requires `@trigger.dev/sdk` to be installed (it's a peer dependency). + * + * @example + * ```ts + * import { chatTask } from "@trigger.dev/ai"; + * import { streamText, convertToModelMessages } from "ai"; + * import { openai } from "@ai-sdk/openai"; + * + * // Simple: return streamText result — auto-piped to the frontend + * export const myChatTask = chatTask({ + * id: "my-chat-task", + * run: async ({ messages }) => { + * return streamText({ + * model: openai("gpt-4o"), + * messages: convertToModelMessages(messages), + * }); + * }, + * }); + * ``` + * + * @example + * ```ts + * import { chatTask, pipeChat } from "@trigger.dev/ai"; + * + * // Complex: use pipeChat() from deep inside your agent code + * export const myAgentTask = chatTask({ + * id: "my-agent-task", + * run: async ({ messages }) => { + * await runComplexAgentLoop(messages); + * // pipeChat() called internally by the agent loop + * }, + * }); + * ``` + */ +export function chatTask( + options: ChatTaskOptions +): Task { + const { run: userRun, ...restOptions } = options; + + return createTask({ + ...restOptions, + run: async (payload: ChatTaskPayload) => { + const result = await userRun(payload); + + // If the run function returned a StreamTextResult or similar, + // automatically pipe it to the chat stream + if (isUIMessageStreamable(result)) { + await pipeChat(result); + } + + return result; + }, + }); +} diff --git a/packages/ai/src/pipeChat.ts b/packages/ai/src/pipeChat.ts new file mode 100644 index 000000000..885951c59 --- /dev/null +++ b/packages/ai/src/pipeChat.ts @@ -0,0 +1,137 @@ +import { realtimeStreams } from "@trigger.dev/core/v3"; + +/** + * The default stream key used for chat transport communication. + * + * Both `TriggerChatTransport` (frontend) and `pipeChat` (backend) use this key + * by default to ensure they communicate over the same stream. + */ +export const CHAT_STREAM_KEY = "chat"; + +/** + * Options for `pipeChat`. + */ +export type PipeChatOptions = { + /** + * Override the stream key to pipe to. + * Must match the `streamKey` option on `TriggerChatTransport`. + * + * @default "chat" + */ + streamKey?: string; + + /** + * An AbortSignal to cancel the stream. + */ + signal?: AbortSignal; + + /** + * The target run ID to pipe the stream to. + * @default "self" (current run) + */ + target?: string; +}; + +/** + * An object that has a `toUIMessageStream()` method, like the result of `streamText()` from the AI SDK. + */ +type UIMessageStreamable = { + toUIMessageStream: (...args: any[]) => AsyncIterable | ReadableStream; +}; + +function isUIMessageStreamable(value: unknown): value is UIMessageStreamable { + return ( + typeof value === "object" && + value !== null && + "toUIMessageStream" in value && + typeof (value as any).toUIMessageStream === "function" + ); +} + +function isAsyncIterable(value: unknown): value is AsyncIterable { + return ( + typeof value === "object" && + value !== null && + Symbol.asyncIterator in value + ); +} + +function isReadableStream(value: unknown): value is ReadableStream { + return ( + typeof value === "object" && + value !== null && + typeof (value as any).getReader === "function" + ); +} + +/** + * Pipes a chat stream to the realtime stream, making it available to the + * `TriggerChatTransport` on the frontend. + * + * Accepts any of: + * - A `StreamTextResult` from the AI SDK (has `.toUIMessageStream()`) + * - An `AsyncIterable` of `UIMessageChunk`s + * - A `ReadableStream` of `UIMessageChunk`s + * + * This must be called from inside a Trigger.dev task's `run` function. + * + * @example + * ```ts + * import { task } from "@trigger.dev/sdk"; + * import { pipeChat, type ChatTaskPayload } from "@trigger.dev/ai"; + * import { streamText, convertToModelMessages } from "ai"; + * + * export const myChatTask = task({ + * id: "my-chat-task", + * run: async (payload: ChatTaskPayload) => { + * const result = streamText({ + * model: openai("gpt-4o"), + * messages: convertToModelMessages(payload.messages), + * }); + * + * await pipeChat(result); + * }, + * }); + * ``` + * + * @example + * ```ts + * // Deep inside your agent library — pipeChat works from anywhere inside a task + * async function runAgentLoop(messages: CoreMessage[]) { + * const result = streamText({ model, messages }); + * await pipeChat(result); + * } + * ``` + * + * @param source - A StreamTextResult, AsyncIterable, or ReadableStream of UIMessageChunks + * @param options - Optional configuration + * @returns A promise that resolves when the stream has been fully piped + */ +export async function pipeChat( + source: UIMessageStreamable | AsyncIterable | ReadableStream, + options?: PipeChatOptions +): Promise { + const streamKey = options?.streamKey ?? CHAT_STREAM_KEY; + + // Resolve the source to an AsyncIterable or ReadableStream + let stream: AsyncIterable | ReadableStream; + + if (isUIMessageStreamable(source)) { + stream = source.toUIMessageStream(); + } else if (isAsyncIterable(source) || isReadableStream(source)) { + stream = source; + } else { + throw new Error( + "pipeChat: source must be a StreamTextResult (with .toUIMessageStream()), " + + "an AsyncIterable, or a ReadableStream" + ); + } + + // Pipe to the realtime stream + const instance = realtimeStreams.pipe(streamKey, stream, { + signal: options?.signal, + target: options?.target, + }); + + await instance.wait(); +} diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 88bf43173..91ae99388 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -8,7 +8,7 @@ export type TriggerChatTransportOptions = { * The Trigger.dev task ID to trigger for chat completions. * This task will receive the chat messages as its payload. */ - taskId: string; + task: string; /** * An access token for authenticating with the Trigger.dev API. @@ -36,8 +36,8 @@ export type TriggerChatTransportOptions = { /** * The stream key where the task pipes UIMessageChunk data. - * Your task must pipe the AI SDK stream to this same key using - * `streams.pipe(streamKey, result.toUIMessageStream())`. + * When using `chatTask()` or `pipeChat()`, this is handled automatically. + * Only set this if you're using a custom stream key. * * @default "chat" */ @@ -59,15 +59,16 @@ export type TriggerChatTransportOptions = { }; /** - * The payload shape that TriggerChatTransport sends to the triggered task. + * The payload shape that the transport sends to the triggered task. * - * Use this type to type your task's `run` function payload: + * When using `chatTask()`, the payload is automatically typed — you don't need + * to import this type. When using `task()` directly, use this type to annotate + * your payload: * * @example * ```ts - * import { task, streams } from "@trigger.dev/sdk"; - * import { streamText, convertToModelMessages } from "ai"; - * import type { ChatTaskPayload } from "@trigger.dev/ai"; + * import { task } from "@trigger.dev/sdk"; + * import { pipeChat, type ChatTaskPayload } from "@trigger.dev/ai"; * * export const myChatTask = task({ * id: "my-chat-task", @@ -76,9 +77,7 @@ export type TriggerChatTransportOptions = { * model: openai("gpt-4o"), * messages: convertToModelMessages(payload.messages), * }); - * - * const { waitUntilComplete } = streams.pipe("chat", result.toUIMessageStream()); - * await waitUntilComplete(); + * await pipeChat(result); * }, * }); * ``` @@ -110,6 +109,7 @@ export type ChatTaskPayload = { /** * Internal state for tracking active chat sessions, used for stream reconnection. + * @internal */ export type ChatSessionState = { runId: string; diff --git a/packages/trigger-sdk/package.json b/packages/trigger-sdk/package.json index f5f840bc1..a3101f703 100644 --- a/packages/trigger-sdk/package.json +++ b/packages/trigger-sdk/package.json @@ -24,7 +24,8 @@ "./package.json": "./package.json", ".": "./src/v3/index.ts", "./v3": "./src/v3/index.ts", - "./ai": "./src/v3/ai.ts" + "./ai": "./src/v3/ai.ts", + "./chat": "./src/v3/chat.ts" }, "sourceDialects": [ "@triggerdotdev/source" @@ -37,6 +38,9 @@ ], "ai": [ "dist/commonjs/v3/ai.d.ts" + ], + "chat": [ + "dist/commonjs/v3/chat.d.ts" ] } }, @@ -123,6 +127,17 @@ "types": "./dist/commonjs/v3/ai.d.ts", "default": "./dist/commonjs/v3/ai.js" } + }, + "./chat": { + "import": { + "@triggerdotdev/source": "./src/v3/chat.ts", + "types": "./dist/esm/v3/chat.d.ts", + "default": "./dist/esm/v3/chat.js" + }, + "require": { + "types": "./dist/commonjs/v3/chat.d.ts", + "default": "./dist/commonjs/v3/chat.js" + } } }, "main": "./dist/commonjs/v3/index.js", diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 59afa2fe2..9e79df22b 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -3,11 +3,16 @@ import { isSchemaZodEsque, Task, type inferSchemaIn, + type PipeStreamOptions, + type TaskOptions, type TaskSchema, type TaskWithSchema, } from "@trigger.dev/core/v3"; +import type { UIMessage } from "ai"; import { dynamicTool, jsonSchema, JSONSchema7, Schema, Tool, ToolCallOptions, zodSchema } from "ai"; import { metadata } from "./metadata.js"; +import { streams } from "./streams.js"; +import { createTask } from "./shared.js"; const METADATA_KEY = "tool.execute.options"; @@ -116,3 +121,240 @@ export const ai = { tool: toolFromTask, currentToolOptions: getToolOptionsFromMetadata, }; + +// --------------------------------------------------------------------------- +// Chat transport helpers — backend side +// --------------------------------------------------------------------------- + +/** + * The default stream key used for chat transport communication. + * Both `TriggerChatTransport` (frontend) and `pipeChat`/`chatTask` (backend) + * use this key by default. + */ +export const CHAT_STREAM_KEY = "chat"; + +/** + * The payload shape that the chat transport sends to the triggered task. + * + * When using `chatTask()`, the payload is automatically typed — you don't need + * to import this type. Use this type only if you're using `task()` directly + * with `pipeChat()`. + */ +export type ChatTaskPayload = { + /** The conversation messages */ + messages: TMessage[]; + + /** The unique identifier for the chat session */ + chatId: string; + + /** + * The trigger type: + * - `"submit-message"`: A new user message + * - `"regenerate-message"`: Regenerate the last assistant response + */ + trigger: "submit-message" | "regenerate-message"; + + /** The ID of the message to regenerate (only for `"regenerate-message"`) */ + messageId?: string; + + /** Custom metadata from the frontend */ + metadata?: unknown; +}; + +/** + * Options for `pipeChat`. + */ +export type PipeChatOptions = { + /** + * Override the stream key. Must match the `streamKey` on `TriggerChatTransport`. + * @default "chat" + */ + streamKey?: string; + + /** An AbortSignal to cancel the stream. */ + signal?: AbortSignal; + + /** + * The target run ID to pipe to. + * @default "self" (current run) + */ + target?: string; +}; + +/** + * An object with a `toUIMessageStream()` method (e.g. `StreamTextResult` from `streamText()`). + */ +type UIMessageStreamable = { + toUIMessageStream: (...args: any[]) => AsyncIterable | ReadableStream; +}; + +function isUIMessageStreamable(value: unknown): value is UIMessageStreamable { + return ( + typeof value === "object" && + value !== null && + "toUIMessageStream" in value && + typeof (value as any).toUIMessageStream === "function" + ); +} + +function isAsyncIterable(value: unknown): value is AsyncIterable { + return typeof value === "object" && value !== null && Symbol.asyncIterator in value; +} + +function isReadableStream(value: unknown): value is ReadableStream { + return typeof value === "object" && value !== null && typeof (value as any).getReader === "function"; +} + +/** + * Pipes a chat stream to the realtime stream, making it available to the + * `TriggerChatTransport` on the frontend. + * + * Accepts: + * - A `StreamTextResult` from `streamText()` (has `.toUIMessageStream()`) + * - An `AsyncIterable` of `UIMessageChunk`s + * - A `ReadableStream` of `UIMessageChunk`s + * + * Must be called from inside a Trigger.dev task's `run` function. + * + * @example + * ```ts + * import { task } from "@trigger.dev/sdk"; + * import { pipeChat, type ChatTaskPayload } from "@trigger.dev/sdk/ai"; + * import { streamText, convertToModelMessages } from "ai"; + * + * export const myChatTask = task({ + * id: "my-chat-task", + * run: async (payload: ChatTaskPayload) => { + * const result = streamText({ + * model: openai("gpt-4o"), + * messages: convertToModelMessages(payload.messages), + * }); + * + * await pipeChat(result); + * }, + * }); + * ``` + * + * @example + * ```ts + * // Works from anywhere inside a task — even deep in your agent code + * async function runAgentLoop(messages: CoreMessage[]) { + * const result = streamText({ model, messages }); + * await pipeChat(result); + * } + * ``` + */ +export async function pipeChat( + source: UIMessageStreamable | AsyncIterable | ReadableStream, + options?: PipeChatOptions +): Promise { + const streamKey = options?.streamKey ?? CHAT_STREAM_KEY; + + let stream: AsyncIterable | ReadableStream; + + if (isUIMessageStreamable(source)) { + stream = source.toUIMessageStream(); + } else if (isAsyncIterable(source) || isReadableStream(source)) { + stream = source; + } else { + throw new Error( + "pipeChat: source must be a StreamTextResult (with .toUIMessageStream()), " + + "an AsyncIterable, or a ReadableStream" + ); + } + + const pipeOptions: PipeStreamOptions = {}; + if (options?.signal) { + pipeOptions.signal = options.signal; + } + if (options?.target) { + pipeOptions.target = options.target; + } + + const { waitUntilComplete } = streams.pipe(streamKey, stream, pipeOptions); + await waitUntilComplete(); +} + +/** + * Options for defining a chat task. + * + * Extends the standard `TaskOptions` but pre-types the payload as `ChatTaskPayload` + * and overrides `run` to accept `ChatTaskPayload` directly. + * + * **Auto-piping:** If the `run` function returns a value with `.toUIMessageStream()` + * (like a `StreamTextResult`), the stream is automatically piped to the frontend. + * For complex flows, use `pipeChat()` manually from anywhere in your code. + */ +export type ChatTaskOptions = Omit< + TaskOptions, + "run" +> & { + /** + * The run function for the chat task. + * + * Receives a `ChatTaskPayload` with the conversation messages, chat session ID, + * and trigger type. + * + * **Auto-piping:** If this function returns a value with `.toUIMessageStream()`, + * the stream is automatically piped to the frontend. + */ + run: (payload: ChatTaskPayload) => Promise; +}; + +/** + * Creates a Trigger.dev task pre-configured for AI SDK chat. + * + * - **Pre-types the payload** as `ChatTaskPayload` — no manual typing needed + * - **Auto-pipes the stream** if `run` returns a `StreamTextResult` + * - For complex flows, use `pipeChat()` from anywhere inside your task code + * + * @example + * ```ts + * import { chatTask } from "@trigger.dev/sdk/ai"; + * import { streamText, convertToModelMessages } from "ai"; + * import { openai } from "@ai-sdk/openai"; + * + * // Simple: return streamText result — auto-piped to the frontend + * export const myChatTask = chatTask({ + * id: "my-chat-task", + * run: async ({ messages }) => { + * return streamText({ + * model: openai("gpt-4o"), + * messages: convertToModelMessages(messages), + * }); + * }, + * }); + * ``` + * + * @example + * ```ts + * import { chatTask, pipeChat } from "@trigger.dev/sdk/ai"; + * + * // Complex: pipeChat() from deep in your agent code + * export const myAgentTask = chatTask({ + * id: "my-agent-task", + * run: async ({ messages }) => { + * await runComplexAgentLoop(messages); + * }, + * }); + * ``` + */ +export function chatTask( + options: ChatTaskOptions +): Task { + const { run: userRun, ...restOptions } = options; + + return createTask({ + ...restOptions, + run: async (payload: ChatTaskPayload) => { + const result = await userRun(payload); + + // Auto-pipe if the run function returned a StreamTextResult or similar + if (isUIMessageStreamable(result)) { + await pipeChat(result); + } + + return result; + }, + }); +} diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts new file mode 100644 index 000000000..5a7872c10 --- /dev/null +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -0,0 +1,294 @@ +/** + * @module @trigger.dev/sdk/chat + * + * Browser-safe module for AI SDK chat transport integration. + * Use this on the frontend with the AI SDK's `useChat` hook. + * + * For backend helpers (`chatTask`, `pipeChat`), use `@trigger.dev/sdk/ai` instead. + * + * @example + * ```tsx + * import { useChat } from "@ai-sdk/react"; + * import { TriggerChatTransport } from "@trigger.dev/sdk/chat"; + * + * function Chat({ accessToken }: { accessToken: string }) { + * const { messages, sendMessage, status } = useChat({ + * transport: new TriggerChatTransport({ + * task: "my-chat-task", + * accessToken, + * }), + * }); + * } + * ``` + */ + +import type { ChatTransport, UIMessage, UIMessageChunk, ChatRequestOptions } from "ai"; +import { ApiClient, SSEStreamSubscription } from "@trigger.dev/core/v3"; + +const DEFAULT_STREAM_KEY = "chat"; +const DEFAULT_BASE_URL = "https://api.trigger.dev"; +const DEFAULT_STREAM_TIMEOUT_SECONDS = 120; + +/** + * Options for creating a TriggerChatTransport. + */ +export type TriggerChatTransportOptions = { + /** + * The Trigger.dev task ID to trigger for chat completions. + * This task should be defined using `chatTask()` from `@trigger.dev/sdk/ai`, + * or a regular `task()` that uses `pipeChat()`. + */ + task: string; + + /** + * An access token for authenticating with the Trigger.dev API. + * + * This must be a token with permission to trigger the task. You can use: + * - A **trigger public token** created via `auth.createTriggerPublicToken(taskId)` (recommended for frontend use) + * - A **secret API key** (for server-side use only — never expose in the browser) + * + * Can also be a function that returns a token string, useful for dynamic token refresh. + */ + accessToken: string | (() => string); + + /** + * Base URL for the Trigger.dev API. + * @default "https://api.trigger.dev" + */ + baseURL?: string; + + /** + * The stream key where the task pipes UIMessageChunk data. + * When using `chatTask()` or `pipeChat()`, this is handled automatically. + * Only set this if you're using a custom stream key. + * + * @default "chat" + */ + streamKey?: string; + + /** + * Additional headers to include in API requests to Trigger.dev. + */ + headers?: Record; + + /** + * The number of seconds to wait for the realtime stream to produce data + * before timing out. + * + * @default 120 + */ + streamTimeoutSeconds?: number; +}; + +/** + * Internal state for tracking active chat sessions. + * @internal + */ +type ChatSessionState = { + runId: string; + publicAccessToken: string; +}; + +/** + * A custom AI SDK `ChatTransport` that runs chat completions as durable Trigger.dev tasks. + * + * When `sendMessages` is called, the transport: + * 1. Triggers a Trigger.dev task with the chat messages as payload + * 2. Subscribes to the task's realtime stream to receive `UIMessageChunk` data + * 3. Returns a `ReadableStream` that the AI SDK processes natively + * + * @example + * ```tsx + * import { useChat } from "@ai-sdk/react"; + * import { TriggerChatTransport } from "@trigger.dev/sdk/chat"; + * + * function Chat({ accessToken }: { accessToken: string }) { + * const { messages, sendMessage, status } = useChat({ + * transport: new TriggerChatTransport({ + * task: "my-chat-task", + * accessToken, + * }), + * }); + * + * // ... render messages + * } + * ``` + * + * On the backend, define the task using `chatTask` from `@trigger.dev/sdk/ai`: + * + * @example + * ```ts + * import { chatTask } from "@trigger.dev/sdk/ai"; + * import { streamText, convertToModelMessages } from "ai"; + * + * export const myChatTask = chatTask({ + * id: "my-chat-task", + * run: async ({ messages }) => { + * return streamText({ + * model: openai("gpt-4o"), + * messages: convertToModelMessages(messages), + * }); + * }, + * }); + * ``` + */ +export class TriggerChatTransport implements ChatTransport { + private readonly taskId: string; + private readonly resolveAccessToken: () => string; + private readonly baseURL: string; + private readonly streamKey: string; + private readonly extraHeaders: Record; + private readonly streamTimeoutSeconds: number; + + private sessions: Map = new Map(); + + constructor(options: TriggerChatTransportOptions) { + this.taskId = options.task; + this.resolveAccessToken = + typeof options.accessToken === "function" + ? options.accessToken + : () => options.accessToken as string; + this.baseURL = options.baseURL ?? DEFAULT_BASE_URL; + this.streamKey = options.streamKey ?? DEFAULT_STREAM_KEY; + this.extraHeaders = options.headers ?? {}; + this.streamTimeoutSeconds = options.streamTimeoutSeconds ?? DEFAULT_STREAM_TIMEOUT_SECONDS; + } + + sendMessages = async ( + options: { + trigger: "submit-message" | "regenerate-message"; + chatId: string; + messageId: string | undefined; + messages: UIMessage[]; + abortSignal: AbortSignal | undefined; + } & ChatRequestOptions + ): Promise> => { + const { trigger, chatId, messageId, messages, abortSignal, body, metadata } = options; + + const payload = { + messages, + chatId, + trigger, + messageId, + metadata, + ...(body ?? {}), + }; + + const currentToken = this.resolveAccessToken(); + const apiClient = new ApiClient(this.baseURL, currentToken); + + const triggerResponse = await apiClient.triggerTask(this.taskId, { + payload: JSON.stringify(payload), + options: { + payloadType: "application/json", + }, + }); + + const runId = triggerResponse.id; + const publicAccessToken = + "publicAccessToken" in triggerResponse + ? (triggerResponse as { publicAccessToken?: string }).publicAccessToken + : undefined; + + this.sessions.set(chatId, { + runId, + publicAccessToken: publicAccessToken ?? currentToken, + }); + + return this.subscribeToStream(runId, publicAccessToken ?? currentToken, abortSignal); + }; + + reconnectToStream = async ( + options: { + chatId: string; + } & ChatRequestOptions + ): Promise | null> => { + const session = this.sessions.get(options.chatId); + if (!session) { + return null; + } + + return this.subscribeToStream(session.runId, session.publicAccessToken, undefined); + }; + + private subscribeToStream( + runId: string, + accessToken: string, + abortSignal: AbortSignal | undefined + ): ReadableStream { + const headers: Record = { + Authorization: `Bearer ${accessToken}`, + ...this.extraHeaders, + }; + + const subscription = new SSEStreamSubscription( + `${this.baseURL}/realtime/v1/streams/${runId}/${this.streamKey}`, + { + headers, + signal: abortSignal, + timeoutInSeconds: this.streamTimeoutSeconds, + } + ); + + return new ReadableStream({ + start: async (controller) => { + try { + const sseStream = await subscription.subscribe(); + const reader = sseStream.getReader(); + + try { + while (true) { + const { done, value } = await reader.read(); + + if (done) { + controller.close(); + return; + } + + if (abortSignal?.aborted) { + reader.cancel(); + reader.releaseLock(); + controller.close(); + return; + } + + controller.enqueue(value.chunk as UIMessageChunk); + } + } catch (readError) { + reader.releaseLock(); + throw readError; + } + } catch (error) { + if (error instanceof Error && error.name === "AbortError") { + controller.close(); + return; + } + + controller.error(error); + } + }, + }); + } +} + +/** + * Creates a new `TriggerChatTransport` instance. + * + * @example + * ```tsx + * import { useChat } from "@ai-sdk/react"; + * import { createChatTransport } from "@trigger.dev/sdk/chat"; + * + * const transport = createChatTransport({ + * task: "my-chat-task", + * accessToken: publicAccessToken, + * }); + * + * function Chat() { + * const { messages, sendMessage } = useChat({ transport }); + * } + * ``` + */ +export function createChatTransport(options: TriggerChatTransportOptions): TriggerChatTransport { + return new TriggerChatTransport(options); +}