The streamable UI has been slow to update. This may be a bug or a performance issue or you forgot to call .done()

[ajaykrupalk]

My first tool call works but the consequent call throws an error

The terminal shows this error The streamable UI has been slow to update. This may be a bug or a performance issue or you forgot to call .done()``

The code for the same is here:

const { value: stream } = await streamUI({
        model: google("gemini-1.5-flash"),
        system: `\
      - you are a friendly assistant
      - reply in lower case
    `,
        messages: messages.get() as CoreMessage[],
        text: async function* ({ content, done }) {
            if (done) {
                messages.done([
                    ...(messages.get() as CoreMessage[]),
                    { role: "assistant", content },
                ]);

                contentStream.done();
            } else {
                contentStream.update(content);
            }

            return textComponent;
        },
        tools: {
            viewUsage: {
                description: "view current usage for youtube",
                parameters: z.object({
                    type: z.enum(["youtube"]),
                }),
                generate: async function* ({ type }) {
                    const toolCallId = generateId();

                    messages.done([
                        ...(messages.get() as CoreMessage[]),
                        {
                            role: "assistant",
                            content: [
                                {
                                    type: "tool-call",
                                    toolCallId,
                                    toolName: "viewUsage",
                                    args: { type },
                                },
                            ],
                        },
                        {
                            role: "tool",
                            content: [
                                {
                                    type: "tool-result",
                                    toolName: "viewUsage",
                                    toolCallId,
                                    result: `The current usage for ${type} is currently displayed on the screen`,
                                },
                            ],
                        },
                    ]);

                    return (
                        <Message role="assistant" content={<UsageView type={type} />} />
                    );
                },
            },
        },
    });

[ajaykrupalk]

@lgrammel would be great to have your help here