Save chat to database on onSetAIState is not working correctly if connection_limit is set to high

[littledino2112]

Description

I'm following the AI chatbot example to develop my own and run into this issue. I'm using supabase's database and hosted the project on vercel. I noticed that the saveChat call sometimes would finish AFTER the new message is updated in aistate. For the chatbot example, It would make the chat_id page refresh and redirect to /chat since when router.refresh is called, the chat message isn't saved to db yet. I notice this happens a lot when I set connection_limit to 2 or 3 in my connection string to supabase.

Code example

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

    const session = await auth()

    if (session && session.user) {
      // Save chat to db if event is 'done'
      // ! Decide not to save the 'user' message since there could be case when it fails to save
      // then all the next messages sent to server will fail due to missing 'user' message
      if (done) {
      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)
}

This is the code in the Chat component where it refreshes the page when messagesLength equals 2

export function Chat({ id, className, session, missingKeys }: ChatProps) {
  const router = useRouter()
  const path = usePathname()
  const [input, setInput] = useState('')
  const [messages] = useUIState()
  const [aiState] = useAIState()

  const [_, setNewChatId] = useLocalStorage('newChatId', id)

  useEffect(() => {
    if (session?.user) {
      if (!path.includes('chat') && messages.length === 1) {
        window.history.replaceState({}, '', `/chat/${id}`)
      }
    }
  }, [id, path, session?.user, messages])

  useEffect(() => {
    const messagesLength = aiState.messages?.length
    if (messagesLength === 2) {
      router.refresh()
    }
  }, [aiState.messages, router])

Additional context

Link to vercel ai chatbot example: https://github.com/vercel/ai-chatbot

[littledino2112]

Someone also mentioned they run into this issue in here.

[Godrules500]

Facing the same issue right now and trying to create a work around

[rishikeshydv]

This sounds like an interesting issue. Since saveChat is an async function, router.refresh() might not wait for the completion of the Promise. Creating an API seems to be a good option. For example: Creating an api router in Nextjs as api/v1/

useEffect (() => {
function handleBeforeUnload {
<your code>
}
if (navigator){
        navigator.sendBeacon("/api/v1/<name>", data);
}
    window.addEventListener("beforeunload", handleBeforeUnload);

    return () => {
      window.removeEventListener("beforeunload", handleBeforeUnload);
    };
},[your dependencies])

[terrytjw]

in prod, my user sometimes can't be found, whereas it works ok on local. weird...

[theitaliandev]

facing the same issue, still trying to find a workaround

[terrytjw]

@theitaliandev instead of using onSetAIState, simply add your saveToDB() logic right after the aiState.done(...)

i commented away the onSetAIState all together. worked for me.

[theitaliandev]

@terrytjw thank you!