From 77a13512762d0b5943191005edb11fe6925116fa Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Mon, 9 Mar 2026 07:51:22 +0000 Subject: [PATCH] feat(chat): add preload support, dynamic tools, and preload-specific timeouts - Add onPreload hook and preloaded field to all lifecycle events - Add transport.preload(chatId) for eagerly starting runs before first message - Add preloadWarmTimeoutInSeconds and preloadTimeout task options - Add preload:true run tag and chat.preloaded span attributes - Add UserTool model for per-user dynamic tools loaded from DB - Load dynamic tools in onPreload/onChatStart via chat.local - Build dynamicTool() instances in run and spread into streamText tools - Reference project: preload on new chat, dynamic company-info and user-preferences tools --- packages/trigger-sdk/src/v3/ai.ts | 164 +++++++++++++++++- packages/trigger-sdk/src/v3/chat.ts | 57 ++++++ references/ai-chat/prisma/schema.prisma | 20 ++- .../ai-chat/src/components/chat-app.tsx | 2 + references/ai-chat/src/trigger/chat.ts | 80 ++++++++- 5 files changed, 310 insertions(+), 13 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 47cad62b5..7ce8c4714 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -300,7 +300,7 @@ const chatStream = streams.define({ id: _CHAT_STREAM_KEY }); type ChatTaskWirePayload = { messages: TMessage[]; chatId: string; - trigger: "submit-message" | "regenerate-message"; + trigger: "submit-message" | "regenerate-message" | "preload"; messageId?: string; metadata?: TMetadata; /** Whether this run is continuing an existing chat whose previous run ended. */ @@ -330,8 +330,9 @@ export type ChatTaskPayload = { * The trigger type: * - `"submit-message"`: A new user message * - `"regenerate-message"`: Regenerate the last assistant response + * - `"preload"`: Run was preloaded before the first message (only on turn 0) */ - trigger: "submit-message" | "regenerate-message"; + trigger: "submit-message" | "regenerate-message" | "preload"; /** The ID of the message to regenerate (only for `"regenerate-message"`) */ messageId?: string; @@ -343,6 +344,8 @@ export type ChatTaskPayload = { continuation: boolean; /** The run ID of the previous run (only set when `continuation` is true). */ previousRunId?: string; + /** Whether this run was preloaded before the first message. */ + preloaded: boolean; }; /** @@ -510,6 +513,20 @@ async function pipeChat( * emits a control chunk and suspends via `messagesInput.wait()`. The frontend * transport resumes the same run by sending the next message via input streams. */ +/** + * Event passed to the `onPreload` callback. + */ +export type PreloadEvent = { + /** The unique identifier for the chat session. */ + chatId: string; + /** The Trigger.dev run ID for this conversation. */ + runId: string; + /** A scoped access token for this chat run. */ + chatAccessToken: string; + /** Custom data from the frontend. */ + clientData?: TClientData; +}; + /** * Event passed to the `onChatStart` callback. */ @@ -528,6 +545,8 @@ export type ChatStartEvent = { continuation: boolean; /** The run ID of the previous run (only set when `continuation` is true). */ previousRunId?: string; + /** Whether this run was preloaded before the first message. */ + preloaded: boolean; }; /** @@ -552,6 +571,8 @@ export type TurnStartEvent = { continuation: boolean; /** The run ID of the previous run (only set when `continuation` is true). */ previousRunId?: string; + /** Whether this run was preloaded before the first message. */ + preloaded: boolean; }; /** @@ -601,6 +622,8 @@ export type TurnCompleteEvent = { continuation: boolean; /** The run ID of the previous run (only set when `continuation` is true). */ previousRunId?: string; + /** Whether this run was preloaded before the first message. */ + preloaded: boolean; }; export type ChatTaskOptions< @@ -638,6 +661,22 @@ export type ChatTaskOptions< */ run: (payload: ChatTaskRunPayload>) => Promise; + /** + * Called when a preloaded run starts, before the first message arrives. + * + * Use this to initialize state, create DB records, and load context early — + * so everything is ready when the user's first message comes through. + * + * @example + * ```ts + * onPreload: async ({ chatId, clientData }) => { + * await db.chat.create({ data: { id: chatId } }); + * userContext.init(await loadUser(clientData.userId)); + * } + * ``` + */ + onPreload?: (event: PreloadEvent>) => Promise | void; + /** * Called on the first turn (turn 0) of a new run, before the `run` function executes. * @@ -722,6 +761,26 @@ export type ChatTaskOptions< * @default "1h" */ chatAccessTokenTTL?: string; + + /** + * How long (in seconds) to keep the run warm after `onPreload` fires, + * waiting for the first message before suspending. + * + * Only applies to preloaded runs (triggered via `transport.preload()`). + * + * @default Same as `warmTimeoutInSeconds` + */ + preloadWarmTimeoutInSeconds?: number; + + /** + * How long to wait (suspended) for the first message after a preloaded run starts. + * If no message arrives within this time, the run ends. + * + * Only applies to preloaded runs. + * + * @default Same as `turnTimeout` + */ + preloadTimeout?: string; }; /** @@ -760,6 +819,7 @@ function chatTask< const { run: userRun, clientDataSchema, + onPreload, onChatStart, onTurnStart, onTurnComplete, @@ -767,6 +827,8 @@ function chatTask< turnTimeout = "1h", warmTimeoutInSeconds = 30, chatAccessTokenTTL = "1h", + preloadWarmTimeoutInSeconds, + preloadTimeout, ...restOptions } = options; @@ -786,6 +848,7 @@ function chatTask< let currentWirePayload = payload; const continuation = payload.continuation ?? false; const previousRunId = payload.previousRunId; + const preloaded = payload.trigger === "preload"; // Accumulated model messages across turns. Turn 1 initialises from the // full history the frontend sends; subsequent turns append only the new @@ -806,6 +869,96 @@ function chatTask< }); try { + // Handle preloaded runs — fire onPreload, then wait for the first real message + if (preloaded) { + if (activeSpan) { + activeSpan.setAttribute("chat.preloaded", true); + } + + const currentRunId = taskContext.ctx?.run.id ?? ""; + let preloadAccessToken = ""; + if (currentRunId) { + try { + preloadAccessToken = await auth.createPublicToken({ + scopes: { + read: { runs: currentRunId }, + write: { inputStreams: currentRunId }, + }, + expirationTime: chatAccessTokenTTL, + }); + } catch { + // Token creation failed + } + } + + // Parse client data for the preload hook + const preloadClientData = (parseClientData + ? await parseClientData(payload.metadata) + : payload.metadata) as inferSchemaOut; + + // Fire onPreload hook + if (onPreload) { + await tracer.startActiveSpan( + "onPreload()", + async () => { + await onPreload({ + chatId: payload.chatId, + runId: currentRunId, + chatAccessToken: preloadAccessToken, + clientData: preloadClientData, + }); + }, + { + attributes: { + [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart", + [SemanticInternalAttributes.COLLAPSED]: true, + "chat.id": payload.chatId, + "chat.preloaded": true, + }, + } + ); + } + + // Wait for the first real message — use preload-specific timeouts if configured + const effectivePreloadWarmTimeout = + (metadata.get(WARM_TIMEOUT_METADATA_KEY) as number | undefined) + ?? preloadWarmTimeoutInSeconds + ?? warmTimeoutInSeconds; + + let firstMessage: ChatTaskWirePayload | undefined; + + if (effectivePreloadWarmTimeout > 0) { + const warm = await messagesInput.once({ + timeoutMs: effectivePreloadWarmTimeout * 1000, + spanName: "preload wait (warm)", + }); + + if (warm.ok) { + firstMessage = warm.output; + } + } + + if (!firstMessage) { + const effectivePreloadTimeout = + (metadata.get(TURN_TIMEOUT_METADATA_KEY) as string | undefined) + ?? preloadTimeout + ?? turnTimeout; + + const suspended = await messagesInput.wait({ + timeout: effectivePreloadTimeout, + spanName: "preload wait (suspended)", + }); + + if (!suspended.ok) { + return; // Timed out waiting for first message — end run + } + + firstMessage = suspended.output; + } + + currentWirePayload = firstMessage; + } + for (let turn = 0; turn < maxTurns; turn++) { // Extract turn-level context before entering the span const { metadata: wireMetadata, messages: uiMessages, ...restWire } = currentWirePayload; @@ -947,6 +1100,7 @@ function chatTask< chatAccessToken: turnAccessToken, continuation, previousRunId, + preloaded, }); }, { @@ -956,6 +1110,7 @@ function chatTask< "chat.id": currentWirePayload.chatId, "chat.messages.count": accumulatedMessages.length, "chat.continuation": continuation, + "chat.preloaded": preloaded, ...(previousRunId ? { "chat.previous_run_id": previousRunId } : {}), }, } @@ -978,6 +1133,7 @@ function chatTask< clientData, continuation, previousRunId, + preloaded, }); }, { @@ -989,6 +1145,7 @@ function chatTask< "chat.messages.count": accumulatedMessages.length, "chat.trigger": currentWirePayload.trigger, "chat.continuation": continuation, + "chat.preloaded": preloaded, ...(previousRunId ? { "chat.previous_run_id": previousRunId } : {}), }, } @@ -1014,6 +1171,7 @@ function chatTask< clientData, continuation, previousRunId, + preloaded, signal: combinedSignal, cancelSignal, stopSignal, @@ -1119,6 +1277,7 @@ function chatTask< stopped: wasStopped, continuation, previousRunId, + preloaded, }); }, { @@ -1129,6 +1288,7 @@ function chatTask< "chat.turn": turn + 1, "chat.stopped": wasStopped, "chat.continuation": continuation, + "chat.preloaded": preloaded, ...(previousRunId ? { "chat.previous_run_id": previousRunId } : {}), "chat.messages.count": accumulatedMessages.length, "chat.response.parts.count": capturedResponseMessage?.parts?.length ?? 0, diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index f8aff0726..bf269c88e 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -183,6 +183,7 @@ export type TriggerChatTransportOptions = { /** Priority (lower = higher priority). */ priority?: number; }; + }; /** @@ -451,6 +452,62 @@ export class TriggerChatTransport implements ChatTransport { this._onSessionChange = callback; } + /** + * Eagerly trigger a run for a chat before the first message is sent. + * This allows initialization (DB setup, context loading) to happen + * while the user is still typing, reducing first-response latency. + * + * The task's `onPreload` hook fires immediately. The run then waits + * for the first message via input stream. When `sendMessages` is called + * later, it detects the existing session and sends via input stream + * instead of triggering a new run. + * + * No-op if a session already exists for this chatId. + */ + async preload(chatId: string): Promise { + // Don't preload if session already exists + if (this.sessions.get(chatId)?.runId) return; + + const payload = { + messages: [] as never[], + chatId, + trigger: "preload" as const, + metadata: this.defaultMetadata, + }; + + const currentToken = await this.resolveAccessToken(); + const apiClient = new ApiClient(this.baseURL, currentToken); + + const autoTags = [`chat:${chatId}`, "preload:true"]; + const userTags = this.triggerOptions?.tags ?? []; + const tags = [...autoTags, ...userTags].slice(0, 5); + + const triggerResponse = await apiClient.triggerTask(this.taskId, { + payload, + options: { + payloadType: "application/json", + tags, + queue: this.triggerOptions?.queue ? { name: this.triggerOptions.queue } : undefined, + maxAttempts: this.triggerOptions?.maxAttempts, + machine: this.triggerOptions?.machine, + priority: this.triggerOptions?.priority, + }, + }); + + const runId = triggerResponse.id; + const publicAccessToken = + "publicAccessToken" in triggerResponse + ? (triggerResponse as { publicAccessToken?: string }).publicAccessToken + : undefined; + + const newSession: ChatSessionState = { + runId, + publicAccessToken: publicAccessToken ?? currentToken, + }; + this.sessions.set(chatId, newSession); + this.notifySessionChange(chatId, newSession); + } + private notifySessionChange( chatId: string, session: ChatSessionState | null diff --git a/references/ai-chat/prisma/schema.prisma b/references/ai-chat/prisma/schema.prisma index d3941b750..4cecb78bc 100644 --- a/references/ai-chat/prisma/schema.prisma +++ b/references/ai-chat/prisma/schema.prisma @@ -8,14 +8,24 @@ datasource db { } model User { - id String @id + id String @id name String - plan String @default("free") // "free" | "pro" + plan String @default("free") // "free" | "pro" preferredModel String? - messageCount Int @default(0) - createdAt DateTime @default(now()) - updatedAt DateTime @updatedAt + messageCount Int @default(0) + createdAt DateTime @default(now()) + updatedAt DateTime @updatedAt chats Chat[] + tools UserTool[] +} + +model UserTool { + id String @id @default(cuid()) + userId String + name String + description String + responseTemplate String @default("") + user User @relation(fields: [userId], references: [id]) } model Chat { diff --git a/references/ai-chat/src/components/chat-app.tsx b/references/ai-chat/src/components/chat-app.tsx index c1008e2fe..e304e2b3f 100644 --- a/references/ai-chat/src/components/chat-app.tsx +++ b/references/ai-chat/src/components/chat-app.tsx @@ -98,6 +98,8 @@ export function ChatApp({ setActiveChatId(id); setMessages([]); setNewChatModel(DEFAULT_MODEL); + // Eagerly start the run — onPreload fires immediately for initialization + transport.preload(id); } function handleSelectChat(id: string) { diff --git a/references/ai-chat/src/trigger/chat.ts b/references/ai-chat/src/trigger/chat.ts index 321f279ba..8611266ee 100644 --- a/references/ai-chat/src/trigger/chat.ts +++ b/references/ai-chat/src/trigger/chat.ts @@ -1,7 +1,7 @@ import { chat, ai } from "@trigger.dev/sdk/ai"; import { schemaTask } from "@trigger.dev/sdk"; -import { streamText, tool, stepCountIs, generateId } from "ai"; -import type { LanguageModel } from "ai"; +import { streamText, tool, dynamicTool, stepCountIs, generateId } from "ai"; +import type { LanguageModel, Tool as AITool } from "ai"; import { openai } from "@ai-sdk/openai"; import { anthropic } from "@ai-sdk/anthropic"; import { z } from "zod"; @@ -136,6 +136,11 @@ const userContext = chat.local<{ messageCount: number; }>(); +// Per-run dynamic tools — loaded from DB in onPreload/onChatStart +const userToolDefs = chat.local< + Array<{ name: string; description: string; responseTemplate: string }> +>(); + // -------------------------------------------------------------------------- // Subtask: deep research — fetches multiple URLs and streams progress // back to the parent chat via chat.stream using data-* chunks @@ -244,8 +249,8 @@ export const aiChat = chat.task({ clientDataSchema: z.object({ model: z.string().optional(), userId: z.string() }), warmTimeoutInSeconds: 60, chatAccessTokenTTL: "2h", - onChatStart: async ({ chatId, runId, chatAccessToken, clientData, continuation }) => { - // Load user context from DB — available for the entire run + onPreload: async ({ chatId, runId, chatAccessToken, clientData }) => { + // Eagerly initialize before the user's first message arrives const user = await prisma.user.upsert({ where: { id: clientData.userId }, create: { id: clientData.userId, name: "User" }, @@ -259,8 +264,57 @@ export const aiChat = chat.task({ messageCount: user.messageCount, }); + // Load user-specific dynamic tools + const tools = await prisma.userTool.findMany({ where: { userId: clientData.userId } }); + userToolDefs.init(tools); + + // Create chat record and session + await prisma.chat.upsert({ + where: { id: chatId }, + create: { + id: chatId, + title: "New chat", + userId: user.id, + model: clientData?.model ?? DEFAULT_MODEL, + }, + update: {}, + }); + await prisma.chatSession.upsert({ + where: { id: chatId }, + create: { id: chatId, runId, publicAccessToken: chatAccessToken }, + update: { runId, publicAccessToken: chatAccessToken }, + }); + }, + onChatStart: async ({ chatId, runId, chatAccessToken, clientData, continuation, preloaded }) => { + if (preloaded) { + // Already initialized in onPreload — just update session + await prisma.chatSession.upsert({ + where: { id: chatId }, + create: { id: chatId, runId, publicAccessToken: chatAccessToken }, + update: { runId, publicAccessToken: chatAccessToken }, + }); + return; + } + + // Non-preloaded path: full initialization + const user = await prisma.user.upsert({ + where: { id: clientData.userId }, + create: { id: clientData.userId, name: "User" }, + update: {}, + }); + userContext.init({ + userId: user.id, + name: user.name, + plan: user.plan as "free" | "pro", + preferredModel: user.preferredModel, + messageCount: user.messageCount, + }); + + // Load user-specific dynamic tools + const tools = await prisma.userTool.findMany({ where: { userId: clientData.userId } }); + userToolDefs.init(tools); + if (!continuation) { - // Brand new chat — create the record with the selected model await prisma.chat.upsert({ where: { id: chatId }, create: { @@ -273,7 +327,6 @@ export const aiChat = chat.task({ }); } - // Always update session for the new run await prisma.chatSession.upsert({ where: { id: chatId }, create: { id: chatId, runId, publicAccessToken: chatAccessToken }, @@ -328,6 +381,20 @@ export const aiChat = chat.task({ const modelId = clientData?.model ?? userContext.preferredModel ?? undefined; const useReasoning = REASONING_MODELS.has(modelId ?? DEFAULT_MODEL); + // Build dynamic tools from user's DB-configured tools (loaded in onPreload/onChatStart) + const dynamicTools: Record> = {}; + for (const t of userToolDefs.value ?? []) { + dynamicTools[t.name] = dynamicTool({ + description: t.description, + inputSchema: z.object({ + query: z.string().describe("The query or topic to look up"), + }), + execute: async (input) => { + return { result: t.responseTemplate.replace("{{query}}", (input as any).query) }; + }, + }); + } + return streamText({ model: getModel(modelId), system: `You are a helpful assistant for ${userContext.name} (${userContext.plan} plan). Be concise and friendly.`, @@ -336,6 +403,7 @@ export const aiChat = chat.task({ inspectEnvironment, webFetch, deepResearch: ai.tool(deepResearch), + ...dynamicTools, }, stopWhen: stepCountIs(10), abortSignal: stopSignal,