github-featured/triggerdotdev--trigger.dev
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

feat(chat): add preload support, dynamic tools, and preload-specific timeouts

Browse Source 
- 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
This commit is contained in:
Eric Allam
2026-03-09 07:51:22 +00:00
parent c98f6f3c8d
commit 77a1351276
5 changed files with 310 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/trigger-sdk/src/v3/ai.ts
+162 -2
View File
@@ -300,7 +300,7 @@ const chatStream = streams.define<UIMessageChunk>({ id: _CHAT_STREAM_KEY });
type ChatTaskWirePayload<TMessage extends UIMessage = UIMessage, TMetadata = unknown> = {
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<TClientData = unknown> = {
* 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<TClientData = unknown> = {
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<TClientData = unknown> = {
/** 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<TClientData = unknown> = {
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<TClientData = unknown> = {
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<TClientData = unknown> = {
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<inferSchemaOut<TClientDataSchema>>) => Promise<unknown>;
/**
* 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<inferSchemaOut<TClientDataSchema>>) => Promise<void> | 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<TClientDataSchema>;
// 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,
packages/trigger-sdk/src/v3/chat.ts
+57
View File
@@ -183,6 +183,7 @@ export type TriggerChatTransportOptions<TClientData = unknown> = {
/** Priority (lower = higher priority). */
priority?: number;
};
};
/**
@@ -451,6 +452,62 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
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<void> {
// 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
references/ai-chat/prisma/schema.prisma
+15 -5
View File
@@ -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 {
references/ai-chat/src/components/chat-app.tsx
+2
View File
@@ -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) {
references/ai-chat/src/trigger/chat.ts
+74 -6
View File
@@ -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<string, AITool<unknown, unknown>> = {};
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,