Skip to content

Fix Chat Page Refreshing Before Saving Chat #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
+124 −45
Open

Fix Chat Page Refreshing Before Saving Chat #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e5af3b1
Resolve race condition when `aiState.done` is called before `onSetAIS…
Saran33 Aug 13, 2024
d236bf7
Change `checkPurchaseStatus` interval from every second to every 5 se…
Saran33 Aug 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion components/chat.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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'
Expand All @@ -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)

Expand All @@ -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])

Expand Down
59 changes: 53 additions & 6 deletions components/stocks/stock-purchase.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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>()
Expand Down Expand Up @@ -59,6 +61,51 @@ export function Purchase({
setAIState({ ...aiState, messages: [...aiState.messages, message] })
}

useEffect(() => {
const checkPurchaseStatus = () => {
if (purchaseStatus !== 'requires_action') {
return;
}
// 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');
}
}
}
Comment on lines +65 to +93
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The logic within checkPurchaseStatus could benefit from additional comments to clarify the order of operations and the conditions under which the purchase status is updated. This would improve readability and maintainability.

};
checkPurchaseStatus();

let intervalId: NodeJS.Timeout | null = null;
if (purchaseStatus === 'requires_action') {
intervalId = setInterval(checkPurchaseStatus, 5000);
}

return () => {
if (intervalId) {
clearInterval(intervalId);
}
};
}, [purchaseStatus, toolCallId, aiState.messages]);
Comment on lines +64 to +107
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The useEffect hook depends on aiState.messages. If setPurchaseStatus triggers a re-render and subsequently modifies aiState.messages (directly or indirectly), it could lead to an infinite loop. While the current logic seems to prevent this, it's a pattern that warrants careful consideration. Consider using a more stable dependency or a ref to store the previous state.


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">
Expand All @@ -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>
Expand Down Expand Up @@ -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>
Expand Down
89 changes: 53 additions & 36 deletions lib/chat/actions.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import 'server-only'

import type { MutableAIState } from '@/lib/types'
import {
createAI,
createStreamableUI,
Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -82,7 +84,7 @@ async function confirmPurchase(symbol: string, price: number, amount: number) {
</SystemMessage>
)

aiState.done({
aiStateDone(aiState, {
...aiState.get(),
messages: [
...aiState.get().messages,
Expand Down Expand Up @@ -159,7 +161,7 @@ async function submitUserMessage(content: string) {

if (done) {
textStream.done()
aiState.done({
aiStateDone(aiState, {
...aiState.get(),
messages: [
...aiState.get().messages,
Expand Down Expand Up @@ -199,7 +201,7 @@ async function submitUserMessage(content: string) {

const toolCallId = nanoid()

aiState.done({
aiStateDone(aiState, {
...aiState.get(),
messages: [
...aiState.get().messages,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -381,6 +382,7 @@ async function submitUserMessage(content: string) {
{
id: nanoid(),
role: 'tool',
createdAt: unixTsNow(),
content: [
{
type: 'tool-result',
Expand All @@ -404,6 +406,7 @@ async function submitUserMessage(content: string) {
numberOfShares,
symbol,
price: +price,
toolCallId,
status: 'requires_action'
}}
/>
Expand Down Expand Up @@ -437,7 +440,7 @@ async function submitUserMessage(content: string) {

const toolCallId = nanoid()

aiState.done({
aiStateDone(aiState, {
...aiState.get(),
messages: [
...aiState.get().messages,
Expand Down Expand Up @@ -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)
Comment on lines +525 to +549
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

⚠️ Potential issue

updateChat omits the chatId field and is fragile when the first message is non-text

  1. The Chat object you persist contains id but not chatId, whereas downstream code (see getUIStateFromAIState) relies on chatId. When the chat is read back from the DB the field will be missing and aiState.chatId will be undefined, breaking UI-key generation and router paths.

  2. title is produced via messages[0].content.substring(...).
    • If messages is empty this will throw.
    • For tool calls or any structured message, content is an array, not a string – calling substring will crash.

A robust fix:

-  const firstMessageContent = messages[0].content as string
-  const title = firstMessageContent.substring(0, 100)
+  const firstContent =
+    typeof messages[0]?.content === 'string' ? messages[0].content : ''
+  const title = firstContent.slice(0, 100)
   const chat: Chat = {
+    chatId,               // keep in sync with UI layer
     id: chatId,
@@
     path
   }
🧰 Tools
🪛 Biome (1.9.4)

[error] 530-530: Change to an optional chain.

Unsafe fix: Change to an optional chain.

(lint/complexity/useOptionalChain)

} 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);
});
};
Comment on lines +555 to +561
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

aiStateDone needs error handling & preserves race-condition fix

Good idea to gate aiState.done behind updateChat, but if updateChat throws, aiState.done is silently skipped and UI never updates. Wrap the await in try / catch / finally so the UI is never blocked by a persistence failure while still logging the error:

-runAsyncFnWithoutBlocking(async () => {
-  await updateChat(newState);
-  aiState.done(newState);
-});
+runAsyncFnWithoutBlocking(async () => {
+  try {
+    await updateChat(newState);
+  } catch (err) {
+    console.error('[aiStateDone] failed to persist chat', err);
+    // TODO: surface non-fatal toast to the user if desired
+  } finally {
+    aiState.done(newState);
+  }
+});
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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);
});
};
const aiStateDone = (aiState: MutableAIState<AIState>, newState: AIState) => {
runAsyncFnWithoutBlocking(async () => {
// resolves race condition in aiState.done - the UI refreshed before db was updated
try {
await updateChat(newState);
} catch (err) {
console.error('[aiStateDone] failed to persist chat', err);
// TODO: surface non-fatal toast to the user if desired
} finally {
aiState.done(newState);
}
});
};


export const getUIStateFromAIState = (aiState: Chat) => {
return aiState.messages
Expand All @@ -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>
Expand Down
7 changes: 7 additions & 0 deletions lib/types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import { CoreMessage } from 'ai'

export type Message = CoreMessage & {
id: string
createdAt?: number;
}

export interface Chat extends Record<string, any> {
Expand Down Expand Up @@ -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);
};
8 changes: 6 additions & 2 deletions lib/utils.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
Comment on lines +54 to +56
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This catch block logs errors that occur within the async function. This is good for debugging. Consider adding more context to the error message, such as the function name or arguments.

Suggested change
fn().catch(error => {
console.error('An error occurred in the async function:', error);
});
fn().catch(error => {
console.error(`An error occurred in the async function ${fn.name}:`, error);
});

};

export const sleep = (ms: number) =>
new Promise(resolve => setTimeout(resolve, ms))
Expand Down Expand Up @@ -87,3 +89,5 @@ export const getMessageFromCode = (resultCode: string) => {
return 'Logged in!'
}
}

export const unixTsNow = () => Math.floor(Date.now() / 1000);