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)