feat: add chat.local for per-run typed data with Proxy access and dirty tracking

Author: ericallam

packages/trigger-sdk/src/v3/ai.ts

@@ -1062,6 +1062,165 @@ function setWarmTimeoutInSeconds(seconds: number): void {
   metadata.set(WARM_TIMEOUT_METADATA_KEY, seconds);
 }
 
+// ---------------------------------------------------------------------------
+// chat.local — per-run typed data with Proxy access
+// ---------------------------------------------------------------------------
+
+/** @internal Symbol for storing the locals key on the proxy target. */
+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;
+
+/**
+ * A Proxy-backed, run-scoped data object that appears as `T` to users.
+ * Includes helper methods for initialization, dirty tracking, and serialization.
+ * Internal metadata is stored behind Symbols and invisible to
+ * `Object.keys()`, `JSON.stringify()`, and spread.
+ */
+export type ChatLocal<T extends Record<string, unknown>> = T & {
+  /** Initialize the local with a value. Call in `onChatStart` or `run()`. */
+  init(value: T): void;
+  /** Returns `true` if any property was set since the last check. Resets the dirty flag. */
+  hasChanged(): boolean;
+  /** Returns a plain object copy of the current value. Useful for persistence. */
+  get(): T;
+  readonly [CHAT_LOCAL_KEY]: ReturnType<typeof locals.create<T>>;
+  readonly [CHAT_LOCAL_DIRTY_KEY]: ReturnType<typeof locals.create<boolean>>;
+};
+
+/**
+ * Creates a per-run typed data object accessible from anywhere during task execution.
+ *
+ * Declare at module level, then initialize inside a lifecycle hook (e.g. `onChatStart`)
+ * using `chat.initLocal()`. Properties are accessible directly via the Proxy.
+ *
+ * Multiple locals can coexist — each gets its own isolated run-scoped storage.
+ *
+ * @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 }>();
+ *
+ * export const myChat = chat.task({
+ *   id: "my-chat",
+ *   onChatStart: async ({ clientData }) => {
+ *     const prefs = await db.prefs.findUnique({ where: { userId: clientData.userId } });
+ *     userPrefs.init(prefs ?? { theme: "dark", language: "en" });
+ *     gameState.init({ score: 0, streak: 0 });
+ *   },
+ *   onTurnComplete: async ({ chatId }) => {
+ *     if (gameState.hasChanged()) {
+ *       await db.save({ where: { chatId }, data: gameState.get() });
+ *     }
+ *   },
+ *   run: async ({ messages }) => {
+ *     gameState.score++;
+ *     return streamText({
+ *       system: `User prefers ${userPrefs.theme} theme. Score: ${gameState.score}`,
+ *       messages,
+ *     });
+ *   },
+ * });
+ * ```
+ */
+function chatLocal<T extends Record<string, unknown>>(): ChatLocal<T> {
+  const localKey = locals.create<T>(`chat.local.${chatLocalCounter++}`);
+  const dirtyKey = locals.create<boolean>(`chat.local.${chatLocalCounter++}.dirty`);
+
+  const target = {} as any;
+  target[CHAT_LOCAL_KEY] = localKey;
+  target[CHAT_LOCAL_DIRTY_KEY] = dirtyKey;
+
+  return new Proxy(target, {
+    get(_target, prop, _receiver) {
+      // Internal Symbol properties
+      if (prop === CHAT_LOCAL_KEY) return _target[CHAT_LOCAL_KEY];
+      if (prop === CHAT_LOCAL_DIRTY_KEY) return _target[CHAT_LOCAL_DIRTY_KEY];
+
+      // Instance methods
+      if (prop === "init") {
+        return (value: T) => {
+          locals.set(localKey, value);
+          locals.set(dirtyKey, false);
+        };
+      }
+      if (prop === "hasChanged") {
+        return () => {
+          const dirty = locals.get(dirtyKey) ?? false;
+          locals.set(dirtyKey, false);
+          return dirty;
+        };
+      }
+      if (prop === "get") {
+        return () => {
+          const current = locals.get(localKey);
+          if (current === undefined) {
+            throw new Error(
+              "local.get() called before initialization. Call local.init() first."
+            );
+          }
+          return { ...current };
+        };
+      }
+      // toJSON for serialization (JSON.stringify(local))
+      if (prop === "toJSON") {
+        return () => {
+          const current = locals.get(localKey);
+          return current ? { ...current } : undefined;
+        };
+      }
+
+      const current = locals.get(localKey);
+      if (current === undefined) return undefined;
+      return (current as any)[prop];
+    },
+
+    set(_target, prop, value) {
+      // Don't allow setting internal Symbols
+      if (typeof prop === "symbol") return false;
+
+      const current = locals.get(localKey);
+      if (current === undefined) {
+        throw new Error(
+          "chat.local can only be modified after initialization. " +
+            "Call local.init() in onChatStart or run() first."
+        );
+      }
+      locals.set(localKey, { ...current, [prop]: value });
+      locals.set(dirtyKey, true);
+      return true;
+    },
+
+    has(_target, prop) {
+      if (typeof prop === "symbol") return prop in _target;
+      const current = locals.get(localKey);
+      return current !== undefined && prop in current;
+    },
+
+    ownKeys() {
+      const current = locals.get(localKey);
+      return current ? Reflect.ownKeys(current) : [];
+    },
+
+    getOwnPropertyDescriptor(_target, prop) {
+      if (typeof prop === "symbol") return undefined;
+      const current = locals.get(localKey);
+      if (current === undefined || !(prop in current)) return undefined;
+      return {
+        configurable: true,
+        enumerable: true,
+        writable: true,
+        value: (current as any)[prop],
+      };
+    },
+  }) as ChatLocal<T>;
+}
+
+
 /**
  * Extracts the client data (metadata) type from a chat task.
  * Use this to type the `metadata` option on the transport.
@@ -1088,6 +1247,8 @@ export const chat = {
   task: chatTask,
   /** Pipe a stream to the chat transport. See {@link pipeChat}. */
   pipe: pipeChat,
+  /** Create a per-run typed local. See {@link chatLocal}. */
+  local: chatLocal,
   /** Create a public access token for a chat task. See {@link createChatAccessToken}. */
   createAccessToken: createChatAccessToken,
   /** Override the turn timeout at runtime (duration string). See {@link setTurnTimeout}. */

references/ai-chat/prisma/migrations/20260306165319_add_user_model/migration.sql

@@ -0,0 +1,18 @@
+-- AlterTable
+ALTER TABLE "Chat" ADD COLUMN     "userId" TEXT;
+
+-- CreateTable
+CREATE TABLE "User" (
+    "id" TEXT NOT NULL,
+    "name" TEXT NOT NULL,
+    "plan" TEXT NOT NULL DEFAULT 'free',
+    "preferredModel" TEXT,
+    "messageCount" INTEGER NOT NULL DEFAULT 0,
+    "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updatedAt" TIMESTAMP(3) NOT NULL,
+
+    CONSTRAINT "User_pkey" PRIMARY KEY ("id")
+);
+
+-- AddForeignKey
+ALTER TABLE "Chat" ADD CONSTRAINT "Chat_userId_fkey" FOREIGN KEY ("userId") REFERENCES "User"("id") ON DELETE SET NULL ON UPDATE CASCADE;

references/ai-chat/prisma/schema.prisma

@@ -7,10 +7,23 @@ datasource db {
   provider = "postgresql"
 }
 
+model User {
+  id             String   @id
+  name           String
+  plan           String   @default("free") // "free" | "pro"
+  preferredModel String?
+  messageCount   Int      @default(0)
+  createdAt      DateTime @default(now())
+  updatedAt      DateTime @updatedAt
+  chats          Chat[]
+}
+
 model Chat {
   id        String   @id
   title     String
   messages  Json     @default("[]")
+  userId    String?
+  user      User?    @relation(fields: [userId], references: [id])
   createdAt DateTime @default(now())
   updatedAt DateTime @updatedAt
 }

references/ai-chat/src/trigger/chat.ts

@@ -84,15 +84,38 @@ const inspectEnvironment = tool({
 declare const Bun: unknown;
 declare const Deno: unknown;
 
+// Per-run user context — loaded from DB in onChatStart, accessible everywhere
+const userContext = chat.local<{
+  userId: string;
+  name: string;
+  plan: "free" | "pro";
+  preferredModel: string | null;
+  messageCount: number;
+}>();
+
 export const aiChat = chat.task({
   id: "ai-chat",
   clientDataSchema: z.object({ model: z.string().optional(), userId: z.string() }),
   warmTimeoutInSeconds: 60,
   chatAccessTokenTTL: "2h",
-  onChatStart: async ({ chatId, runId, chatAccessToken }) => {
+  onChatStart: async ({ chatId, runId, chatAccessToken, clientData }) => {
+    // Load user context from DB — available for the entire run
+    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,
+    });
+
     await prisma.chat.upsert({
       where: { id: chatId },
-      create: { id: chatId, title: "New chat" },
+      create: { id: chatId, title: "New chat", userId: user.id },
       update: {},
     });
     await prisma.chatSession.upsert({
@@ -113,7 +136,7 @@ export const aiChat = chat.task({
       update: { runId, publicAccessToken: chatAccessToken },
     });
   },
-  onTurnComplete: async ({ chatId, uiMessages, runId, chatAccessToken, lastEventId }) => {
+  onTurnComplete: async ({ chatId, uiMessages, runId, chatAccessToken, lastEventId, clientData }) => {
     // Persist final messages + assistant response + stream position
     await prisma.chat.update({
       where: { id: chatId },
@@ -124,11 +147,33 @@ export const aiChat = chat.task({
       create: { id: chatId, runId, publicAccessToken: chatAccessToken, lastEventId },
       update: { runId, publicAccessToken: chatAccessToken, lastEventId },
     });
+
+    // Persist user context changes (message count, preferred model) if anything changed
+    if (userContext.hasChanged()) {
+      await prisma.user.update({
+        where: { id: userContext.userId },
+        data: {
+          messageCount: userContext.messageCount,
+          preferredModel: userContext.preferredModel,
+        },
+      });
+    }
   },
   run: async ({ messages, clientData, stopSignal }) => {
+    // Track usage
+    userContext.messageCount++;
+
+    // Remember their model choice
+    if (clientData?.model) {
+      userContext.preferredModel = clientData.model;
+    }
+
+    // Use preferred model if none specified
+    const modelId = clientData?.model ?? userContext.preferredModel ?? undefined;
+
     return streamText({
-      model: getModel(clientData?.model),
-      system: "You are a helpful assistant. Be concise and friendly.",
+      model: getModel(modelId),
+      system: `You are a helpful assistant for ${userContext.name} (${userContext.plan} plan). Be concise and friendly.`,
       messages,
       tools: { inspectEnvironment },
       stopWhen: stepCountIs(10),