From 7ef73da7bd40569cb2dd067eaf861d875422f986 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Mon, 9 Mar 2026 11:08:48 +0000 Subject: [PATCH] feat(chat): auto-hydrate chat.local values in ai.tool subtasks --- packages/trigger-sdk/src/v3/ai.ts | 95 ++++++++++++++++++++++---- references/ai-chat/src/trigger/chat.ts | 4 +- 2 files changed, 84 insertions(+), 15 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 7ce8c4714..b4717def0 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -40,6 +40,8 @@ export type ToolCallExecutionOptions = { turn?: number; continuation?: boolean; clientData?: unknown; + /** Serialized chat.local values from the parent run. @internal */ + chatLocals?: Record; }; /** Chat context stored in locals during each chat.task turn for auto-detection. */ @@ -121,6 +123,18 @@ function toolFromTask< toolMeta.clientData = chatCtx.clientData; } + // Serialize initialized chat.local values for subtask hydration + const chatLocals: Record = {}; + for (const entry of chatLocalRegistry) { + const value = locals.get(entry.key); + if (value !== undefined) { + chatLocals[entry.id] = value; + } + } + if (Object.keys(chatLocals).length > 0) { + toolMeta.chatLocals = chatLocals; + } + return await task .triggerAndWait(input as inferSchemaIn, { metadata: { @@ -1546,8 +1560,31 @@ function cleanupAbortedParts(message: UIMessage): UIMessage { const CHAT_LOCAL_KEY: unique symbol = Symbol("chatLocalKey"); /** @internal Symbol for storing the dirty-tracking locals key. */ const CHAT_LOCAL_DIRTY_KEY: unique symbol = Symbol("chatLocalDirtyKey"); -/** @internal Counter for generating unique locals IDs. */ -let chatLocalCounter = 0; + +// --------------------------------------------------------------------------- +// chat.local registry — tracks all declared locals for serialization +// --------------------------------------------------------------------------- + +type ChatLocalEntry = { key: ReturnType; id: string }; +const chatLocalRegistry = new Set(); + +/** @internal Run-scoped flag to ensure hydration happens at most once per run. */ +const chatLocalsHydratedKey = locals.create("chat.locals.hydrated"); + +/** + * Hydrate chat.local values from subtask metadata (set by toolFromTask). + * Runs once per run — subsequent calls are no-ops. + * @internal + */ +function hydrateLocalsFromMetadata(): void { + if (locals.get(chatLocalsHydratedKey)) return; + locals.set(chatLocalsHydratedKey, true); + const opts = metadata.get(METADATA_KEY) as ToolCallExecutionOptions | undefined; + if (!opts?.chatLocals) return; + for (const [id, value] of Object.entries(opts.chatLocals)) { + locals.set(locals.create(id), value); + } +} /** * A Proxy-backed, run-scoped data object that appears as `T` to users. @@ -1574,12 +1611,16 @@ export type ChatLocal> = T & { * * Multiple locals can coexist — each gets its own isolated run-scoped storage. * + * The `id` is required and must be unique across all `chat.local()` calls in + * your project. It's used to serialize values into subtask metadata so that + * `ai.tool()` subtasks can auto-hydrate parent locals (read-only). + * * @example * ```ts * import { chat } from "@trigger.dev/sdk/ai"; * - * const userPrefs = chat.local<{ theme: string; language: string }>(); - * const gameState = chat.local<{ score: number; streak: number }>(); + * const userPrefs = chat.local<{ theme: string; language: string }>({ id: "userPrefs" }); + * const gameState = chat.local<{ score: number; streak: number }>({ id: "gameState" }); * * export const myChat = chat.task({ * id: "my-chat", @@ -1603,9 +1644,12 @@ export type ChatLocal> = T & { * }); * ``` */ -function chatLocal>(): ChatLocal { - const localKey = locals.create(`chat.local.${chatLocalCounter++}`); - const dirtyKey = locals.create(`chat.local.${chatLocalCounter++}.dirty`); +function chatLocal>(options: { id: string }): ChatLocal { + const id = `chat.local.${options.id}`; + const localKey = locals.create(id); + const dirtyKey = locals.create(`${id}.dirty`); + + chatLocalRegistry.add({ key: localKey, id }); const target = {} as any; target[CHAT_LOCAL_KEY] = localKey; @@ -1633,7 +1677,11 @@ function chatLocal>(): ChatLocal { } if (prop === "get") { return () => { - const current = locals.get(localKey); + let current = locals.get(localKey); + if (current === undefined) { + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } if (current === undefined) { throw new Error( "local.get() called before initialization. Call local.init() first." @@ -1645,12 +1693,21 @@ function chatLocal>(): ChatLocal { // toJSON for serialization (JSON.stringify(local)) if (prop === "toJSON") { return () => { - const current = locals.get(localKey); + let current = locals.get(localKey); + if (current === undefined) { + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } return current ? { ...current } : undefined; }; } - const current = locals.get(localKey); + let current = locals.get(localKey); + if (current === undefined) { + // Auto-hydrate from parent metadata in subtask context + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } if (current === undefined) return undefined; return (current as any)[prop]; }, @@ -1673,18 +1730,30 @@ function chatLocal>(): ChatLocal { has(_target, prop) { if (typeof prop === "symbol") return prop in _target; - const current = locals.get(localKey); + let current = locals.get(localKey); + if (current === undefined) { + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } return current !== undefined && prop in current; }, ownKeys() { - const current = locals.get(localKey); + let current = locals.get(localKey); + if (current === undefined) { + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } return current ? Reflect.ownKeys(current) : []; }, getOwnPropertyDescriptor(_target, prop) { if (typeof prop === "symbol") return undefined; - const current = locals.get(localKey); + let current = locals.get(localKey); + if (current === undefined) { + hydrateLocalsFromMetadata(); + current = locals.get(localKey); + } if (current === undefined || !(prop in current)) return undefined; return { configurable: true, diff --git a/references/ai-chat/src/trigger/chat.ts b/references/ai-chat/src/trigger/chat.ts index 8611266ee..52ba1865d 100644 --- a/references/ai-chat/src/trigger/chat.ts +++ b/references/ai-chat/src/trigger/chat.ts @@ -134,12 +134,12 @@ const userContext = chat.local<{ plan: "free" | "pro"; preferredModel: string | null; messageCount: number; -}>(); +}>({ id: "userContext" }); // Per-run dynamic tools — loaded from DB in onPreload/onChatStart const userToolDefs = chat.local< Array<{ name: string; description: string; responseTemplate: string }> ->(); +>({ id: "userToolDefs" }); // -------------------------------------------------------------------------- // Subtask: deep research — fetches multiple URLs and streams progress