When using streamUI, I need to be able to update the ui from onFinish

[Godrules500]

Description

In the onFinish of streamUI it returns certain values that need to modify how the UI is updated. For instance, if there is a content-filter issue, I need to display to the user of the issue.

If I don't finish the stream in the text callback, and I try and close it in the onFinish, I get UI stream is already closed.

Shouldn't it be that I finish the stream in onFinish?

Also, onFinish is being called before text: isDone= true

Code example

Here are 2 ways that are causing the issue. I continue to get throw new Error(method + ': UI stream is already closed.'); even if I remove all .done

        text: async ({ content, done, delta }) => {
          text += delta || ''
          try {
            let index = getMessageIndex(aiState)
            if (done) {
              reply.update(<BotMessage isFinal={done} index={index} content={text}></BotMessage>)
            } else {
              reply.update(<BotMessage isFinal={done} index={index} content={text}></BotMessage>)
              // textStream.update(delta)
            }
            return <BotMessage isFinal={done} index={index} content={text}></BotMessage>
          } catch (error) {
            console.error(error)
          }
        },
        onFinish: async finishResult => {
          reply.done(
            <BotMessage isFinal={true} index={index}>
              {text}
            </BotMessage>
          )

text: async ({ content, done, delta }) => {
          text += delta || ''
          try {
            let index = getMessageIndex(aiState)
            if (!textStream) {
              textStream = createStreamableValue('')
              textNode = <BotMessage content={textStream.value} isFinal={done} index={index} />
            }
            if (done) {
              textStream.done()
              textNode = <BotMessage isFinal={done} index={index} content={content}></BotMessage>
            } else {
              textStream.update(delta)
            }
            return textNode
          } catch (error) {
            console.error(error)
          }
        },
        onFinish: async finishResult => {
          let done = true

          let title = await titlePromise
          let { error, showReportAnIssue } = getFinishedErrorText(finishResult.finishReason, text)
          if (error) {
            let message = aiState.get().chat.messages[aiState.get().chat.messages.length - 1]
            loggerWithFields({ user: user, message: message.content, promptViolation: true }).info('PROMPT_VIOLATION')
            await handleMessageError(aiState, reply, new Error(error), showReportAnIssue, title)
            return
          }

          if (done) {
            streamFinished(user, text, userFiles, title)
          }
        }

Additional context

No response

[lgrammel]

Can you try streamText with createStreamableValue / createStreamableUI? This will give you more control over that aspect.