From 90a9556e89656c24845fe3eb316cbe52ce64f323 Mon Sep 17 00:00:00 2001
From: Eric Allam <eallam@icloud.com>
Date: Sat, 7 Mar 2026 15:37:57 +0000
Subject: [PATCH] Some improvements to the example ai-chat

---
 references/ai-chat/prisma/schema.prisma       |   1 +
 references/ai-chat/src/app/actions.ts         |   3 +-
 .../ai-chat/src/components/chat-app.tsx       |  47 +++-
 references/ai-chat/src/components/chat.tsx    | 209 +++++++++++++-----
 references/ai-chat/src/trigger/chat.ts        |   9 +-
 5 files changed, 198 insertions(+), 71 deletions(-)

diff --git a/references/ai-chat/prisma/schema.prisma b/references/ai-chat/prisma/schema.prisma
index 5b58955c2..d3941b750 100644
--- a/references/ai-chat/prisma/schema.prisma
+++ b/references/ai-chat/prisma/schema.prisma
@@ -21,6 +21,7 @@ model User {
 model Chat {
   id        String   @id
   title     String
+  model     String   @default("gpt-4o-mini")
   messages  Json     @default("[]")
   userId    String?
   user      User?    @relation(fields: [userId], references: [id])
diff --git a/references/ai-chat/src/app/actions.ts b/references/ai-chat/src/app/actions.ts
index 3b6c55e71..56398c9c8 100644
--- a/references/ai-chat/src/app/actions.ts
+++ b/references/ai-chat/src/app/actions.ts
@@ -8,12 +8,13 @@ export const getChatToken = async () => chat.createAccessToken<typeof aiChat>("a
 
 export async function getChatList() {
   const chats = await prisma.chat.findMany({
-    select: { id: true, title: true, createdAt: true, updatedAt: true },
+    select: { id: true, title: true, model: true, createdAt: true, updatedAt: true },
     orderBy: { updatedAt: "desc" },
   });
   return chats.map((c) => ({
     id: c.id,
     title: c.title,
+    model: c.model,
     createdAt: c.createdAt.getTime(),
     updatedAt: c.updatedAt.getTime(),
   }));
diff --git a/references/ai-chat/src/components/chat-app.tsx b/references/ai-chat/src/components/chat-app.tsx
index a00695ec4..c1008e2fe 100644
--- a/references/ai-chat/src/components/chat-app.tsx
+++ b/references/ai-chat/src/components/chat-app.tsx
@@ -7,6 +7,7 @@ import type { aiChat } from "@/trigger/chat";
 import { useCallback, useEffect, useState } from "react";
 import { Chat } from "@/components/chat";
 import { ChatSidebar } from "@/components/chat-sidebar";
+import { DEFAULT_MODEL } from "@/lib/models";
 import {
   getChatToken,
   getChatList,
@@ -19,18 +20,22 @@ import {
 type ChatMeta = {
   id: string;
   title: string;
+  model: string;
   createdAt: number;
   updatedAt: number;
 };
 
+type SessionInfo = {
+  runId: string;
+  publicAccessToken: string;
+  lastEventId?: string;
+};
+
 type ChatAppProps = {
   initialChatList: ChatMeta[];
   initialActiveChatId: string | null;
   initialMessages: UIMessage[];
-  initialSessions: Record<
-    string,
-    { runId: string; publicAccessToken: string; lastEventId?: string }
-  >;
+  initialSessions: Record<string, SessionInfo>;
 };
 
 export function ChatApp({
@@ -42,15 +47,21 @@ export function ChatApp({
   const [chatList, setChatList] = useState<ChatMeta[]>(initialChatList);
   const [activeChatId, setActiveChatId] = useState<string | null>(initialActiveChatId);
   const [messages, setMessages] = useState<UIMessage[]>(initialMessages);
+  const [sessions, setSessions] = useState<Record<string, SessionInfo>>(initialSessions);
+
+  // Model for new chats (before first message is sent)
+  const [newChatModel, setNewChatModel] = useState(DEFAULT_MODEL);
 
   const handleSessionChange = useCallback(
-    (
-      chatId: string,
-      session: { runId: string; publicAccessToken: string; lastEventId?: string } | null
-    ) => {
-      // Session creation and token updates are handled server-side via onChatStart/onTurnComplete.
-      // We only need to clean up when the run ends (session = null).
-      if (!session) {
+    (chatId: string, session: SessionInfo | null) => {
+      if (session) {
+        setSessions((prev) => ({ ...prev, [chatId]: session }));
+      } else {
+        setSessions((prev) => {
+          const next = { ...prev };
+          delete next[chatId];
+          return next;
+        });
         deleteSessionAction(chatId);
       }
     },
@@ -86,6 +97,7 @@ export function ChatApp({
     const id = generateId();
     setActiveChatId(id);
     setMessages([]);
+    setNewChatModel(DEFAULT_MODEL);
   }
 
   function handleSelectChat(id: string) {
@@ -119,6 +131,14 @@ export function ChatApp({
     setChatList(list);
   }, []);
 
+  // Determine the model for the active chat
+  const activeChatMeta = chatList.find((c) => c.id === activeChatId);
+  const isNewChat = activeChatId != null && !activeChatMeta;
+  const activeModel = isNewChat ? newChatModel : (activeChatMeta?.model ?? DEFAULT_MODEL);
+
+  // Get session for the active chat
+  const activeSession = activeChatId ? sessions[activeChatId] : undefined;
+
   return (
     <main className="flex h-screen">
       <ChatSidebar
@@ -136,6 +156,11 @@ export function ChatApp({
             initialMessages={messages}
             transport={transport}
             resume={messages.length > 0}
+            model={activeModel}
+            isNewChat={isNewChat}
+            onModelChange={isNewChat ? setNewChatModel : undefined}
+            session={activeSession}
+            dashboardUrl={process.env.NEXT_PUBLIC_TRIGGER_DASHBOARD_URL}
             onFirstMessage={handleFirstMessage}
             onMessagesChange={handleMessagesChange}
           />
diff --git a/references/ai-chat/src/components/chat.tsx b/references/ai-chat/src/components/chat.tsx
index 13abe0df6..f6e1916b5 100644
--- a/references/ai-chat/src/components/chat.tsx
+++ b/references/ai-chat/src/components/chat.tsx
@@ -5,7 +5,7 @@ import { useChat } from "@ai-sdk/react";
 import type { TriggerChatTransport } from "@trigger.dev/sdk/chat";
 import { useEffect, useRef, useState } from "react";
 import { Streamdown } from "streamdown";
-import { MODEL_OPTIONS, DEFAULT_MODEL } from "@/lib/models";
+import { MODEL_OPTIONS } from "@/lib/models";
 
 function ToolInvocation({ part }: { part: any }) {
   const [expanded, setExpanded] = useState(false);
@@ -70,11 +70,112 @@ function ToolInvocation({ part }: { part: any }) {
   );
 }
 
+function DebugPanel({
+  chatId,
+  model,
+  status,
+  session,
+  dashboardUrl,
+  messageCount,
+}: {
+  chatId: string;
+  model: string;
+  status: string;
+  session?: { runId: string; publicAccessToken: string; lastEventId?: string };
+  dashboardUrl?: string;
+  messageCount: number;
+}) {
+  const [open, setOpen] = useState(false);
+
+  const runUrl =
+    session?.runId && dashboardUrl
+      ? `${dashboardUrl}/runs/${session.runId}`
+      : undefined;
+
+  return (
+    <div className="shrink-0 border-t border-gray-200 bg-gray-50 text-xs text-gray-500">
+      <button
+        type="button"
+        onClick={() => setOpen(!open)}
+        className="flex w-full items-center gap-2 px-4 py-1.5 hover:bg-gray-100"
+      >
+        <span className="font-medium">Debug</span>
+        <span
+          className={`inline-block h-1.5 w-1.5 rounded-full ${
+            status === "streaming"
+              ? "bg-green-500"
+              : session?.runId
+                ? "bg-yellow-500"
+                : "bg-gray-300"
+          }`}
+        />
+        <span>{status}</span>
+        {session?.runId && (
+          <span className="font-mono">{session.runId.slice(0, 16)}...</span>
+        )}
+        <span className="ml-auto text-gray-400">{open ? "▲" : "▼"}</span>
+      </button>
+
+      {open && (
+        <div className="border-t border-gray-200 px-4 py-2 space-y-1">
+          <Row label="Chat ID" value={chatId} mono />
+          <Row label="Model" value={model} />
+          <Row label="Status" value={status} />
+          <Row label="Messages" value={String(messageCount)} />
+          {session ? (
+            <>
+              <Row label="Run ID" value={session.runId} mono link={runUrl} />
+              <Row label="Last Event ID" value={session.lastEventId ?? "—"} mono />
+            </>
+          ) : (
+            <Row label="Session" value="none" />
+          )}
+        </div>
+      )}
+    </div>
+  );
+}
+
+function Row({
+  label,
+  value,
+  mono,
+  link,
+}: {
+  label: string;
+  value: string;
+  mono?: boolean;
+  link?: string;
+}) {
+  return (
+    <div className="flex items-center gap-2">
+      <span className="w-24 shrink-0 text-gray-400">{label}</span>
+      {link ? (
+        <a
+          href={link}
+          target="_blank"
+          rel="noopener noreferrer"
+          className={`truncate text-blue-600 underline ${mono ? "font-mono" : ""}`}
+        >
+          {value}
+        </a>
+      ) : (
+        <span className={`truncate ${mono ? "font-mono" : ""}`}>{value}</span>
+      )}
+    </div>
+  );
+}
+
 type ChatProps = {
   chatId: string;
   initialMessages: UIMessage[];
   transport: TriggerChatTransport;
   resume?: boolean;
+  model: string;
+  isNewChat: boolean;
+  onModelChange?: (model: string) => void;
+  session?: { runId: string; publicAccessToken: string; lastEventId?: string };
+  dashboardUrl?: string;
   onFirstMessage?: (chatId: string, text: string) => void;
   onMessagesChange?: (chatId: string, messages: UIMessage[]) => void;
 };
@@ -84,12 +185,15 @@ export function Chat({
   initialMessages,
   transport,
   resume: resumeProp,
+  model,
+  isNewChat,
+  onModelChange,
+  session,
+  dashboardUrl,
   onFirstMessage,
   onMessagesChange,
 }: ChatProps) {
   const [input, setInput] = useState("");
-  const [model, setModel] = useState(DEFAULT_MODEL);
-  const modelByUserMsgId = useRef<Map<string, string>>(new Map());
   const hasCalledFirstMessage = useRef(false);
 
   const { messages, sendMessage, stop, status, error } = useChat({
@@ -114,7 +218,7 @@ export function Chat({
   }, [messages, chatId, onFirstMessage]);
 
   // Pending message to send after the current turn completes
-  const [pendingMessage, setPendingMessage] = useState<{ text: string; model: string } | null>(null);
+  const [pendingMessage, setPendingMessage] = useState<string | null>(null);
 
   // Handle turn completion: persist messages and auto-send pending message
   const prevStatus = useRef(status);
@@ -124,47 +228,48 @@ export function Chat({
 
     if (!turnCompleted) return;
 
-    // Persist messages when a turn completes — this ensures the final assistant
-    // message content is saved (not the empty placeholder from mid-stream).
+    // Persist messages when a turn completes
     if (messages.length > 0) {
       onMessagesChange?.(chatId, messages);
     }
 
     // Auto-send the pending message
     if (pendingMessage) {
-      const { text, model: pendingMsgModel } = pendingMessage;
+      const text = pendingMessage;
       setPendingMessage(null);
-      pendingModel.current = pendingMsgModel;
-      sendMessage({ text }, { metadata: { model: pendingMsgModel } });
+      sendMessage({ text }, { metadata: { model } });
     }
-  }, [status, messages, chatId, onMessagesChange, sendMessage, pendingMessage]);
-
-  function getModelForAssistantAt(index: number): string | undefined {
-    for (let i = index - 1; i >= 0; i--) {
-      if (messages[i]?.role === "user") {
-        return modelByUserMsgId.current.get(messages[i].id);
-      }
-    }
-    return undefined;
-  }
-
-  const originalSendMessage = sendMessage;
-  function trackedSendMessage(msg: Parameters<typeof sendMessage>[0], opts?: Parameters<typeof sendMessage>[1]) {
-    pendingModel.current = model;
-    originalSendMessage(msg, opts);
-  }
-  const pendingModel = useRef<string>(model);
-
-  const trackedUserIds = useRef<Set<string>>(new Set());
-  for (const msg of messages) {
-    if (msg.role === "user" && !trackedUserIds.current.has(msg.id)) {
-      trackedUserIds.current.add(msg.id);
-      modelByUserMsgId.current.set(msg.id, pendingModel.current);
-    }
-  }
+  }, [status, messages, chatId, onMessagesChange, sendMessage, pendingMessage, model]);
 
   return (
     <div className="flex h-full flex-col bg-white">
+      {/* Model selector for new chats */}
+      {isNewChat && messages.length === 0 && onModelChange && (
+        <div className="shrink-0 border-b border-gray-200 bg-gray-50 px-4 py-2 flex items-center gap-2">
+          <span className="text-xs text-gray-500">Model:</span>
+          <select
+            value={model}
+            onChange={(e) => onModelChange(e.target.value)}
+            className="rounded-md border border-gray-300 px-2 py-1 text-xs text-gray-600 outline-none focus:border-blue-500 focus:ring-1 focus:ring-blue-500"
+          >
+            {MODEL_OPTIONS.map((m) => (
+              <option key={m} value={m}>
+                {m}
+              </option>
+            ))}
+          </select>
+        </div>
+      )}
+
+      {/* Model badge for existing chats */}
+      {(!isNewChat || messages.length > 0) && (
+        <div className="shrink-0 border-b border-gray-200 bg-gray-50 px-4 py-2">
+          <span className="rounded bg-gray-200 px-1.5 py-0.5 text-[10px] font-medium text-gray-500">
+            {model}
+          </span>
+        </div>
+      )}
+
       {/* Messages */}
       <div className="flex-1 space-y-4 overflow-y-auto p-4">
         {messages.length === 0 && (
@@ -177,13 +282,6 @@ export function Chat({
             className={`flex ${message.role === "user" ? "justify-end" : "justify-start"}`}
           >
             <div className={`max-w-[80%] ${message.role === "user" ? "" : "w-full"}`}>
-              {message.role === "assistant" && (
-                <div className="mb-1 flex items-center gap-2 text-[10px] text-gray-400">
-                  <span className="rounded bg-gray-200 px-1.5 py-0.5 font-medium text-gray-500">
-                    {getModelForAssistantAt(messageIndex) ?? DEFAULT_MODEL}
-                  </span>
-                </div>
-              )}
               <div
                 className={`rounded-lg px-4 py-2 text-sm ${
                   message.role === "user"
@@ -246,7 +344,7 @@ export function Chat({
           <div className="flex justify-end">
             <div className="max-w-[80%]">
               <div className="rounded-lg bg-blue-600 px-4 py-2 text-sm text-white opacity-60">
-                {pendingMessage.text}
+                {pendingMessage}
               </div>
               <div className="mt-1 text-right text-[10px] text-gray-400">
                 Queued — will send when current response finishes
@@ -262,14 +360,24 @@ export function Chat({
         </div>
       )}
 
+      {/* Debug panel */}
+      <DebugPanel
+        chatId={chatId}
+        model={model}
+        status={status}
+        session={session}
+        dashboardUrl={dashboardUrl}
+        messageCount={messages.length}
+      />
+
       <form
         onSubmit={(e) => {
           e.preventDefault();
           if (!input.trim()) return;
           if (status === "streaming") {
-            setPendingMessage({ text: input, model });
+            setPendingMessage(input);
           } else {
-            trackedSendMessage({ text: input }, { metadata: { model } });
+            sendMessage({ text: input }, { metadata: { model } });
           }
           setInput("");
         }}
@@ -300,19 +408,6 @@ export function Chat({
             </button>
           )}
         </div>
-        <div className="mt-2">
-          <select
-            value={model}
-            onChange={(e) => setModel(e.target.value)}
-            className="rounded-md border border-gray-300 px-2 py-1 text-xs text-gray-600 outline-none focus:border-blue-500 focus:ring-1 focus:ring-blue-500"
-          >
-            {MODEL_OPTIONS.map((m) => (
-              <option key={m} value={m}>
-                {m}
-              </option>
-            ))}
-          </select>
-        </div>
       </form>
     </div>
   );
diff --git a/references/ai-chat/src/trigger/chat.ts b/references/ai-chat/src/trigger/chat.ts
index d8e171c31..226e6bcad 100644
--- a/references/ai-chat/src/trigger/chat.ts
+++ b/references/ai-chat/src/trigger/chat.ts
@@ -156,10 +156,15 @@ export const aiChat = chat.task({
     });
 
     if (!continuation) {
-      // Brand new chat — create the record
+      // Brand new chat — create the record with the selected model
       await prisma.chat.upsert({
         where: { id: chatId },
-        create: { id: chatId, title: "New chat", userId: user.id },
+        create: {
+          id: chatId,
+          title: "New chat",
+          userId: user.id,
+          model: clientData.model ?? DEFAULT_MODEL,
+        },
         update: {},
       });
     }