Skip to content
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.

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

Page refresh firing before saveChat #302

Open
Godrules500 opened this issue Apr 5, 2024 · 17 comments · May be fixed by #399 or #404
Open

Page refresh firing before saveChat #302

Godrules500 opened this issue Apr 5, 2024 · 17 comments · May be fixed by #399 or #404

Comments

@Godrules500
Copy link

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
Copy link

amxv commented Apr 8, 2024

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

@Godrules500
Copy link
Author

Godrules500 commented Apr 9, 2024

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
Copy link

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

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

@audvin
Copy link

audvin commented Apr 11, 2024

Experiencing the same issue.

@Godrules500
Copy link
Author

@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
Copy link

@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
Copy link

Hi, any news on this issue ?

@navkuun
Copy link

navkuun commented May 23, 2024

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
Copy link

same for me.

@athrael-soju
Copy link

athrael-soju commented May 23, 2024

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
Copy link

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
Copy link

navkuun commented Jul 5, 2024

any fix to this?

@yorickvanzweeden
Copy link

yorickvanzweeden commented Jul 16, 2024

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
Copy link
Author

Godrules500 commented Jul 17, 2024

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 Saran33 linked a pull request Aug 13, 2024 that will close this issue
@Saran33
Copy link

Saran33 commented Aug 13, 2024

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

@Saran33
Copy link

Saran33 commented Aug 13, 2024

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
Copy link

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
10 participants