Text response is lost when `streamUI` calls a tool

[thiagozf]

Description

OpenAI models can generate text content and call tools simultaneously, in a single request. The model usually uses this to explain what it will do next - e.g. "Sure, give me a second to check San Francisco weather.".

This situation causes a weird behavior on streamUI, where the message sent by AI disappears once the tool's generate function finishes executing.

While debugging, it seems the problem is here. We receive text-delta and use the the text renderer to render the message. However, when we receive a tool-call here, the tool's generate renderer is called, completely ignoring the previously received text content.

I expected both text (with done = true) and the tool's generate renderers to be called on the last iteration.

Code example

async function continueConversation(content: string) {
  "use server";

  const aiState = getMutableAIState<typeof AI>();

  aiState.update({
    ...aiState.get(),
    messages: [
      ...aiState.get().messages,
      { id: generateId(), role: "user", content },
    ],
  });
  
  let textStream: undefined | ReturnType<typeof createStreamableValue<string>>;
  let textNode: undefined | React.ReactNode;

  const result = await streamUI({
    model: openai("gpt-4-turbo", { parallelToolCalls: true }),
    system: `\
    You are a helpful virtual assistant.
    
    IMPORTANT: Before invoking any functions, first you must explain the user what you're going to do. At the very end of your response, you execute function calls.
    `,
    messages: aiState.get().messages,
    text: ({ content, done, delta }) => {
      if (!textStream) {
        textStream = createStreamableValue("");
        textNode = <Message role="assistant" content={textStream.value} />;
      }

      if (done) {
        textStream.done();
        aiState.done({
          ...aiState.get(),
          messages: [
            ...aiState.get().messages,
            {
              id: generateId(),
              role: "assistant",
              content,
            },
          ],
        });
      } else {
        textStream.update(delta);
      }

      return textNode;
    },
    tools: {
      getWeather: {
        description: 'Get the weather in a location',
        parameters: z.object({
          location: z.string().describe('The location to get the weather for'),
        }),
        generate: async function (args, tool) {
          const result = { temperature: 72 + Math.floor(Math.random() * 21) - 10 };
          aiState.done({
            ...aiState.get(),
            messages: [
              ...aiState.get().messages,
              ...generateToolMessages(
                tool,
                args,
                result,
              ),
            ],
          });
          return <Weather result={result} />;
        },
      },
    },
  });

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

Additional context

No response