Resolve race condition when aiState.done is called before onSetAIState - prevent the UI from refreshing before the db is updated

Author: Saran33

components/chat.tsx

@@ -1,5 +1,7 @@
 'use client'
 
+import type { AI } from '@/lib/chat/actions'
+
 import { cn } from '@/lib/utils'
 import { ChatList } from '@/components/chat-list'
 import { ChatPanel } from '@/components/chat-panel'
@@ -24,7 +26,7 @@ export function Chat({ id, className, session, missingKeys }: ChatProps) {
   const path = usePathname()
   const [input, setInput] = useState('')
   const [messages] = useUIState()
-  const [aiState] = useAIState()
+  const [aiState] = useAIState<typeof AI>()
 
   const [_, setNewChatId] = useLocalStorage('newChatId', id)
 
@@ -40,6 +42,8 @@ export function Chat({ id, className, session, missingKeys }: ChatProps) {
     const messagesLength = aiState.messages?.length
     if (messagesLength === 2) {
       router.refresh()
+    } else if (messagesLength === 3 && aiState.messages[2].role === 'tool') {
+      router.refresh()
     }
   }, [aiState.messages, router])
 

components/stocks/stock-purchase.tsx

@@ -1,24 +1,26 @@
 'use client'
 
-import { useId, useState } from 'react'
+import { useEffect, useId, useState } from 'react';
 import { useActions, useAIState, useUIState } from 'ai/rsc'
-import { formatNumber } from '@/lib/utils'
+import { formatNumber, unixTsNow } from '@/lib/utils'
 
 import type { AI } from '@/lib/chat/actions'
 
 interface Purchase {
   numberOfShares?: number
   symbol: string
   price: number
+  toolCallId: string
   status: 'requires_action' | 'completed' | 'expired'
 }
 
 export function Purchase({
-  props: { numberOfShares, symbol, price, status = 'expired' }
+  props: { numberOfShares, symbol, price, toolCallId, status = 'requires_action' }
 }: {
   props: Purchase
 }) {
   const [value, setValue] = useState(numberOfShares || 100)
+  const [purchaseStatus, setPurchaseStatus] = useState(status);
   const [purchasingUI, setPurchasingUI] = useState<null | React.ReactNode>(null)
   const [aiState, setAIState] = useAIState<typeof AI>()
   const [, setMessages] = useUIState<typeof AI>()
@@ -59,6 +61,51 @@ export function Purchase({
     setAIState({ ...aiState, messages: [...aiState.messages, message] })
   }
 
+  useEffect(() => {
+    const checkPurchaseStatus = () => {
+      if (purchaseStatus === 'requires_action') {
+        // check for purchase completion
+        // Find the tool message with the matching toolCallId
+        const toolMessage = aiState.messages.find(
+          message => 
+            message.role === 'tool' && 
+            message.content.some(part => part.toolCallId === toolCallId)
+        );
+
+        if (toolMessage) {
+          const toolMessageIndex = aiState.messages.indexOf(toolMessage);
+          // Check if the next message is a system message containing "purchased"
+          const nextMessage = aiState.messages[toolMessageIndex + 1];
+          if (
+            nextMessage?.role === 'system' &&
+            nextMessage.content.includes('purchased')
+          ) {
+            setPurchaseStatus('completed');
+          } else {
+            // Check for expiration
+            const requestedAt = toolMessage.createdAt;
+            if (!requestedAt || unixTsNow() - requestedAt > 30) {
+              setPurchaseStatus('expired');
+            }
+          }
+        }
+      }
+    };
+
+    checkPurchaseStatus();
+
+    let intervalId: NodeJS.Timeout | null = null;
+    if (purchaseStatus === 'requires_action') {
+      intervalId = setInterval(checkPurchaseStatus, 1000);
+    }
+
+    return () => {
+      if (intervalId) {
+        clearInterval(intervalId);
+      }
+    };
+  }, [purchaseStatus, toolCallId, aiState.messages]);
+
   return (
     <div className="p-4 text-green-400 border rounded-xl bg-zinc-950">
       <div className="inline-block float-right px-2 py-1 text-xs rounded-full bg-white/10">
@@ -68,7 +115,7 @@ export function Purchase({
       <div className="text-3xl font-bold">${price}</div>
       {purchasingUI ? (
         <div className="mt-4 text-zinc-200">{purchasingUI}</div>
-      ) : status === 'requires_action' ? (
+      ) : purchaseStatus === 'requires_action' ? (
         <>
           <div className="relative pb-6 mt-6">
             <p>Shares to purchase</p>
@@ -133,12 +180,12 @@ export function Purchase({
             Purchase
           </button>
         </>
-      ) : status === 'completed' ? (
+      ) : purchaseStatus === 'completed' ? (
         <p className="mb-2 text-white">
           You have successfully purchased {value} ${symbol}. Total cost:{' '}
           {formatNumber(value * price)}
         </p>
-      ) : status === 'expired' ? (
+      ) : purchaseStatus === 'expired' ? (
         <p className="mb-2 text-white">Your checkout session has expired!</p>
       ) : null}
     </div>

lib/chat/actions.tsx

@@ -1,5 +1,6 @@
 import 'server-only'
 
+import type { MutableAIState } from '@/lib/types'
 import {
   createAI,
   createStreamableUI,
@@ -29,7 +30,8 @@ import {
   formatNumber,
   runAsyncFnWithoutBlocking,
   sleep,
-  nanoid
+  nanoid,
+  unixTsNow
 } from '@/lib/utils'
 import { saveChat } from '@/app/actions'
 import { SpinnerMessage, UserMessage } from '@/components/stocks/message'
@@ -82,7 +84,7 @@ async function confirmPurchase(symbol: string, price: number, amount: number) {
       </SystemMessage>
     )
 
-    aiState.done({
+    aiStateDone(aiState, {
       ...aiState.get(),
       messages: [
         ...aiState.get().messages,
@@ -159,7 +161,7 @@ async function submitUserMessage(content: string) {
 
       if (done) {
         textStream.done()
-        aiState.done({
+        aiStateDone(aiState, {
           ...aiState.get(),
           messages: [
             ...aiState.get().messages,
@@ -199,7 +201,7 @@ async function submitUserMessage(content: string) {
 
           const toolCallId = nanoid()
 
-          aiState.done({
+          aiStateDone(aiState, {
             ...aiState.get(),
             messages: [
               ...aiState.get().messages,
@@ -259,8 +261,7 @@ async function submitUserMessage(content: string) {
           await sleep(1000)
 
           const toolCallId = nanoid()
-
-          aiState.done({
+          aiStateDone(aiState, {
             ...aiState.get(),
             messages: [
               ...aiState.get().messages,
@@ -319,7 +320,7 @@ async function submitUserMessage(content: string) {
           const toolCallId = nanoid()
 
           if (numberOfShares <= 0 || numberOfShares > 1000) {
-            aiState.done({
+            aiStateDone(aiState, {
               ...aiState.get(),
               messages: [
                 ...aiState.get().messages,
@@ -362,7 +363,7 @@ async function submitUserMessage(content: string) {
 
             return <BotMessage content={'Invalid amount'} />
           } else {
-            aiState.done({
+            aiStateDone(aiState, {
               ...aiState.get(),
               messages: [
                 ...aiState.get().messages,
@@ -381,6 +382,7 @@ async function submitUserMessage(content: string) {
                 {
                   id: nanoid(),
                   role: 'tool',
+                  createdAt: unixTsNow(),
                   content: [
                     {
                       type: 'tool-result',
@@ -404,6 +406,7 @@ async function submitUserMessage(content: string) {
                     numberOfShares,
                     symbol,
                     price: +price,
+                    toolCallId,
                     status: 'requires_action'
                   }}
                 />
@@ -437,7 +440,7 @@ async function submitUserMessage(content: string) {
 
           const toolCallId = nanoid()
 
-          aiState.done({
+          aiStateDone(aiState, {
             ...aiState.get(),
             messages: [
               ...aiState.get().messages,
@@ -517,36 +520,45 @@ export const AI = createAI<AIState, UIState>({
       return
     }
   },
-  onSetAIState: async ({ state }) => {
-    'use server'
+})
 
-    const session = await auth()
+const updateChat = async (state: AIState) => {
+  'use server'
 
-    if (session && session.user) {
-      const { chatId, messages } = state
-
-      const createdAt = new Date()
-      const userId = session.user.id as string
-      const path = `/chat/${chatId}`
-
-      const firstMessageContent = messages[0].content as string
-      const title = firstMessageContent.substring(0, 100)
-
-      const chat: Chat = {
-        id: chatId,
-        title,
-        userId,
-        createdAt,
-        messages,
-        path
-      }
+  const session = await auth()
 
-      await saveChat(chat)
-    } else {
-      return
+  if (session && session.user) {
+    const { chatId, messages } = state
+
+    const createdAt = new Date()
+    const userId = session.user.id as string
+    const path = `/chat/${chatId}`
+
+    const firstMessageContent = messages[0].content as string
+    const title = firstMessageContent.substring(0, 100)
+
+    const chat: Chat = {
+      id: chatId,
+      title,
+      userId,
+      createdAt,
+      messages,
+      path
     }
+
+    await saveChat(chat)
+  } else {
+    return
   }
-})
+};
+
+const aiStateDone = (aiState: MutableAIState<AIState>, newState: AIState) => {
+  runAsyncFnWithoutBlocking(async () => {
+    // resolves race condition in aiState.done - the UI refreshed before db was updated
+    await updateChat(newState);
+    aiState.done(newState);
+  });
+};
 
 export const getUIStateFromAIState = (aiState: Chat) => {
   return aiState.messages
@@ -569,8 +581,13 @@ export const getUIStateFromAIState = (aiState: Chat) => {
               </BotCard>
             ) : tool.toolName === 'showStockPurchase' ? (
               <BotCard>
-                {/* @ts-expect-error */}
-                <Purchase props={tool.result} />
+                <Purchase
+                  // @ts-expect-error
+                  props={{
+                    ...(tool.result as object),
+                    toolCallId: tool.toolCallId,
+                  }}
+                />
               </BotCard>
             ) : tool.toolName === 'getEvents' ? (
               <BotCard>

lib/types.ts

@@ -2,6 +2,7 @@ import { CoreMessage } from 'ai'
 
 export type Message = CoreMessage & {
   id: string
+  createdAt?: number;
 }
 
 export interface Chat extends Record<string, any> {
@@ -39,3 +40,9 @@ export interface User extends Record<string, any> {
   password: string
   salt: string
 }
+
+export type MutableAIState<AIState> = {
+  get: () => AIState;
+  update: (newState: AIState | ((current: AIState) => AIState)) => void;
+  done: ((newState: AIState) => void) | (() => void);
+};

lib/utils.ts

@@ -51,8 +51,10 @@ export const formatNumber = (value: number) =>
 export const runAsyncFnWithoutBlocking = (
   fn: (...args: any) => Promise<any>
 ) => {
-  fn()
-}
+  fn().catch(error => {
+    console.error('An error occurred in the async function:', error);
+  });
+};
 
 export const sleep = (ms: number) =>
   new Promise(resolve => setTimeout(resolve, ms))
@@ -87,3 +89,5 @@ export const getMessageFromCode = (resultCode: string) => {
       return 'Logged in!'
   }
 }
+
+export const unixTsNow = () => Math.floor(Date.now() / 1000);