Bug in `streamUI` and `createStreamableUI`

[ashersamuel8]

Description

I'm using ai/rsc for my application and there seems to be an issue with streamUI(). I am using streamUI() in function (server actions function) called submitUserMessage(...). The function returns the result.value (the UI component to be streamed) AGAIN after the stream is closed/streamUI() is done executing (check the example below). To dig deeper, I wrote the same function using streamText(), inside which I used createStreamableUI() to stream the UI component. What I noticed is that when I close the stream of createStreamableUI() using .done(), the UI component, for some reason, is returned again, leading to other undesirable behavior like component remounting, layout shifts, and animation restarts on the client side. Is there a way to fix/work around this? I don't notice the same behavior when I close the stream with createStreambleValue().

It is also worth noting that my function and UI components work as desired when I don't closecreateStreableUI() using .done(). However, it throws an error in production environment since the connection isn't closed.

Code example

Code using streamUI():

async function submitUserMessage(content: string) {
  "use server"

  const aiState = getMutableAIState<typeof AI>()

  aiState.update({
    ...aiState.get(),
    messages: [
      ...aiState.get().messages,
      {
        id: nanoid(),
        role: "user",
        content,
      },
    ],
  })

  let textStream: undefined | ReturnType<typeof createStreamableValue<string>>
  let textNode: undefined | React.ReactNode

  const result = await streamUI({
    model: google("models/gemini-1.5-flash-latest"),
    initial: <SpinnerMessage />,
    system: system_instructions,
    messages: [
      ...aiState.get().messages.map((message: any) => ({
        role: message.role,
        content: message.content,
        name: message.name,
      })),
    ],
    text: ({ content, done, delta }) => {
      if (!textStream) {
        textStream = createStreamableValue("")
        textNode = <BotMessage content={textStream.value} />
      }

      if (done) {
        textStream.done()
        aiState.done({
          ...aiState.get(),
          messages: [
            ...aiState.get().messages,
            {
              id: nanoid(),
              role: "assistant",
              content,
            },
          ],
        })
      } else {
        textStream.update(delta)
      }
      return textNode
    },
    onFinish: result => {
      console.log(result.finishReason)
    },
  })

  return {
    id: nanoid(),
    display: result.value,
  }
}

Code using streamText() without closing the createStrarmableUI() stream (works as expected on localhost but throws errors in production because the stream isn't closed:

async function submitUserMessage(content: string) {
  "use server"

  const aiState = getMutableAIState<typeof AI>()

  aiState.update({
    ...aiState.get(),
    messages: [
      ...aiState.get().messages,
      {
        id: nanoid(),
        role: "user",
        content,
      },
    ],
  })

  console.log(aiState.get().messages)

  let textStream: undefined | ReturnType<typeof createStreamableValue<string>>
  const spinnerStream = createStreamableUI(<SpinnerMessage />)

  textStream = createStreamableValue("")
  let textNode = createStreamableUI(<SpinnerMessage />)

  ;(async () => {
    const result = streamText({
      model: google("models/gemini-1.5-flash-latest"),
      system: system_instructions,
      messages: [
        ...aiState.get().messages.map((message: any) => ({
          role: message.role,
          content: message.content,
        })),
      ],
    })

    let textContent = ""
    spinnerStream.done(null)

    textNode.update(<BotMessage content={textStream.value} />)

    for await (const delta of (await result).fullStream) {
      const { type } = delta

      if (type === "text-delta") {
        const { textDelta } = delta

        textContent += textDelta
        textStream.update(textDelta)
      }
    }

    textStream.done()
    aiState.done({
      ...aiState.get(),
      messages: [
        ...aiState.get().messages,
        {
          id: nanoid(),
          role: "assistant",
          content: textContent,
        },
      ],
    })
  })()

Additional context

What I'm using

"ai": "^3.2.16",
"next": "14.2.4",
"react": "^18",
"react-dom": "^18",

[Theonlyhamstertoh]

^ Experiencing the same issue where the component remounts again and causes a major flicker and layout shift in the AI generated text.

[Theonlyhamstertoh]

Try to see if you can downgrade to ai@3.2.5 and see if there flicker disappears?