Page refresh firing before saveChat

[Godrules500]

I have a situation every now and then when I create a new chat that, and the saveChat takes a little bit of time, the chat.tsx router.refresh() fires too soon. How can I ensure that it waits for saveChat to finish?

This is the code that is firing before the saveChat finishes. So the assistant responds, and the saveChat finishes after the aiState.done is called.
useEffect(() => { const messagesLength = aiState.messages?.length if (messagesLength === 2) { router.refresh() } }, [aiState.messages, router])

[amxv]

having the same issue atm, were you able to fix it?

[Godrules500]

Hey @amxv yes, I was able to resolve it. Since I'm using gemini, I had to follow the AI SDK 3 example "next-ai-rsc". But essentially it should be the same thing.

Changes I made:

• Added a "saved" property to AIState (to prevent saving multiple times).
• Refactored saveChat called inside of unstable_onSetAIState to a method
• Called await saveChat when text completion "isFinal" was true.

I now call saveChat inside of unstable_onSetAIState like this

unstable_onSetAIState: async ({ state, done }) => {
    'use server'

    const session = await auth()

    if (session && session.user) {

      // only save it once it's done. No need to save it after each streamed result
      if (done && !state.saved) {
        await addOrUpdateChat(state, session)
        state.saved = true
      }
    } else {
      return
    }
  }
})

I then call await saveChat(state) when isFinal is true and set saved to true.

let state = aiState.get()
        state.messages = [
          ...state.messages,
          {
            id: nanoid(),
            role: 'assistant',
            content
          }
        ]

        await saveChat(state)

        reply.done()
        aiState.done({ ...state, done: isFinal, saved: true })

And then the saveChat() is this

async function saveChat(state: AIState, session: Session | null = null) {
  if (!session) {
    session = await auth()
    if (!session?.user?.id) {
      throw new Error('User not Authorized to save')
    }
  }

  const { chatId, messages } = state

  const createdAt = Date.now()
  const userId = session.user?.id as string
  const path = `/chat/${chatId}`
  const title = messages[0].content.substring(0, 100)

  const chat: Chat = {
    id: chatId,
    title,
    userId,
    createdAt,
    path,
    messages
  }

  await dynamoDbService.saveChat(chat)
}

Let me know if that fixes it for you too or if you run into any issues with it please!

[athrael-soju]

@Godrules500 could you please raise a PR with the fix?

Fixes such as this one are super helpful for keeping the template updated.

[audvin]

Experiencing the same issue.

[Godrules500]

@athrael-soju I will try to get to it soon!

I also found another solution to another issue if it helps y'all.
If you go to new chat, and then clear history, and then send a prompt, it will refresh because the chat ids are out of sync. My fix for me is as follow.

on chat.tsx, this is the problem area for me.

useEffect(() => {
    setNewChatId(id)
  })

I updated mine to

  useEffect(() => {
    // This is here, because it resolve a bug when you go to 'new chat' --> clear history, and then generate a new chat, the IDs were out of sync.
    updateAIState({ ...aiState, chatId: id! })
    setNewChatId(id)
  })

then on actions.tsx, I added this.

async function updateAIState(currentAIState: any) {
  'use server'
  const aiState = getMutableAIState<typeof AI>()
  aiState.done({ ...currentAIState })
}

don't forget to add it where submitUserMessage is added.

Then update chat.tsx to

[athrael-soju]

@athrael-soju I will try to get to it soon!

I also found another solution to another issue if it helps y'all. If you go to new chat, and then clear history, and then send a prompt, it will refresh because the chat ids are out of sync. My fix for me is as follow.

on chat.tsx, this is the problem area for me.

useEffect(() => {
    setNewChatId(id)
  })

I updated mine to

  useEffect(() => {
    // This is here, because it resolve a bug when you go to 'new chat' --> clear history, and then generate a new chat, the IDs were out of sync.
    updateAIState({ ...aiState, chatId: id! })
    setNewChatId(id)
  })

then on actions.tsx, I added this.

async function updateAIState(currentAIState: any) {
  'use server'
  const aiState = getMutableAIState<typeof AI>()
  aiState.done({ ...currentAIState })
}

don't forget to add it where submitUserMessage is added.

Then update chat.tsx to

I think you maybe sharing code from the ai-rsc template here. Is this an issue with this template? I tried reproducing this issue and I'm not seeing it.

[yaberkane05]

Hi, any news on this issue ?

[navkuun]

Made the changes as said here, but still facing this issue. Anyone else?

My current code:

export type AIState = {
  chatId: string
  messages: Message[]
  saved?: boolean
}

export type UIState = {
  id: string
  display: React.ReactNode
}[]


async function addOrUpdateChat(state: AIState, session: Session | null = null) {
  if (!session) {
    session = await auth()
    if (!session?.user?.id) {
      throw new Error('User not Authorized to save')
    }
  }

  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 flowchart = await getFlowchart(chatId)

  const chat: Chat = {
    id: chatId,
    title,
    userId,
    createdAt,
    path,
    messages,
    flowchart: flowchart || { nodes: [], edges: [] } // Use existing flowchart data if available, otherwise initialize with empty arrays
  }

  await saveChat(chat)
}

export const AI = createAI<AIState, UIState>({
  actions: {
    submitUserMessage
  },
  initialUIState: [],
  initialAIState: { chatId: nanoid(), messages: [] },
  onGetUIState: async () => {
    'use server'

    const session = await auth()

    if (session && session.user) {
      const aiState = getAIState()

      if (aiState) {
        const uiState = getUIStateFromAIState(aiState)
        return uiState
      }
    } else {
      return
    }
  },
  onSetAIState: async ({ state, done }) => {
    'use server'

    const session = await auth()

    if (session && session.user) {
      // only save it once it's done. No need to save it after each streamed result
      if (done && !state.saved) {
        await addOrUpdateChat(state, session)
        state.saved = true
      }
    } else {
      return
    }
  }
})

[yaberkane05]

same for me.

[athrael-soju]

I'd wait a bit for the next version of the template to come out, as this one still has some breaking issues. They've already closed down several PRs with fixes, which hopefully points to a release soon.

[sudeepkudari0]

I solved this by doing these changes to "/" and "/new" . In "/" route we will check for existing chats and if present we will push to "/chat/[id]", If no chats then push to "/new".

import { ObjectId } from 'bson'
import { db } from '@/lib/db'
import { auth } from '@/auth'
import { getMissingKeys } from '@/app/actions'
import { AI } from '@/lib/chat/actions'
import { Chat } from '@/components/chat'
import { Session } from '@/lib/types'
import { redirect } from 'next/navigation'

export const metadata = {
  title: 'Sudeep Kudari'
}

const getChatsCustom = async (userId: string) => {
  return await db.chat.findMany({
    where: { userId },
    include: {
      messages: true
    },
    orderBy: { createdAt: 'desc' }
  })

}

const IndexPage = async () => {
  const session = (await auth()) as Session
  const missingKeys = await getMissingKeys()

  if (session?.user) {
    const newData = await getChatsCustom(session.user.id)
    console.log(newData)
    if (newData[0]?.messages?.length > 0) {
      return redirect(`/chat/${newData[0].id}`)
    } else {
      return redirect('/new')
    }
  }

  const chatId = new ObjectId().toHexString()
  return (
    <AI initialAIState={{ chatId, messages: [] }}>
      <Chat id={chatId} session={session} missingKeys={missingKeys} />
    </AI>
  )
}

export default IndexPage

And in "/new"

import { ObjectId } from 'bson'
import { db } from '@/lib/db'
import { auth } from '@/auth'
import { getMissingKeys } from '@/app/actions'
import { Session } from '@/lib/types'
import { redirect } from 'next/navigation'

const createChatWithDynamicPath = async (userId: string) => {
  const chatId = new ObjectId().toHexString()
  const path = `/chat/${chatId}`
  return await db.chat.create({
    data: {
      id: chatId,
      userId,
      title: 'New Chat',
      path
    }
  })
}

export default async function NewPage() {

  const session = (await auth()) as Session

  if (session?.user) {
    const newData = await createChatWithDynamicPath(session.user.id)
    return redirect(`/chat/${newData.id}`)
  }
}

[navkuun]

any fix to this?

[yorickvanzweeden]

TL;DR: Do not call aiState.done multiple times for the same action

What I figured out, is the following:

• Client and server correspond using deltas
• State in the server is stored using AsyncLocalStorage
• aiState.update updates the current state and calls onSetAIState. (No updates are sent yet)
• aiState.done does the same as aiState.update AND it computes a diff that will be sent.

If you call aiState.done multiple times, only the first diff will be received by the client. In my case, the AIState was not correctly in sync and the refresh occurred too fast (as subsequent updates were missed)

[Godrules500]

I don't remember what all changes I've made up to this point, but the change I just made has made a huge difference and solved 3 issues.

1. After the second (or greater) prompt response from AI, if you navigate away, and then back, it would remove the last chat.
2. When doing a new chat, instead of staying or going to the new chat, it would go back to the new chat window.
3. If the db save took too long, the page would go back to the home screen.

It seems that the AIState gets out of sync, and these 2 changes to chat.tsx seems to force aiState to get back in sync.

in chat.tsx I have this.

// This is used to fix Vercel ai/rsc bug where the state is not correctly maintained/updated
 //    causing the aiState to get out of sync.
 // It shows up when you navigate away from the current route and then then back to the same route quickly.
 // This invalidates the cache and refreshes the AI states
 useEffect(() => {
   router.refresh()
 }, [router])

 useEffect(() => {
   // Do not update the aiState and local key chat Id if its the same. If this is removed, it was creating several new chatIds, getting out of sync when saving.
   if (aiState.chat.chatId === chatId) return
   // Update the aiState chatId when going to a new chat. If not, it will save with the wrong chatId.
   aiState.chat.chatId = chatId!
   setAIState(aiState)

   // Update the local storage chat Id.
   setNewChatId(chatId)
 }, [aiState, chatId, setAIState, setNewChatId])

I also had to move saveChat out of onSetAIState. I've changed the naming and extracted it to its own method, but this forced the code to wait for the db to update before marking the state as done.

const aiStateDone = async (aiState: MutableAIState<typeof AI>, newMessage: Message, properties?: Partial<AIState>) => {
  const messages = []
  if (newMessage) {
    messages.push(newMessage)
  }
  let state = {
    ...aiState.get(),
    ...properties,
    chat: {
      ...aiState.get().chat,
      messages: [...aiState.get().chat.messages, ...messages]
    }
  }

  // Doing it here to resolve a bug in aiStateDone where if addOrUpdate took a while, the UI refreshed...
  //    before the db was updated.
  await addOrUpdateChat(state)

  await aiState.done(state)
}

Please let me know if this fixes your issue!

[Saran33]

@Godrules500 nice solution, thank you🚀 I incorporated some of those changes into a PR here #399

[Saran33]

This doesn't appear to be an issue on the demo site though so it may have already been resolved by the Vercel team but the changes not applied to this repo. I said I'd add the changes here anyway in any case it's useful to anyone.

[preshetin]

Thanks @Saran33 - I used the code from your #399 PR and it looks like it does fix that!