AI State not updated during iterations

[fran-cadenas-fu]

Description

Hello,

I'm currently trying to use the new streamUI feature to create a chat that keeps asking questions on a certain topic and displays Kick response buttons to de user through tools.

I can't get it to work because the AIState keeps getting reseted to the one in the first iteration. I'm not sure that this is a problem with the functions memoization or something.

Code example

async function submitUserMessage(userInput: string): Promise<UIMessage> {
  "use server";

  console.log("submitUserMessage", userInput);
  const history = getMutableAIState<typeof AI>();

  const ui = await streamUI({
    model: openai("gpt-4-turbo"),
    system: systemPrompt,
    messages: history.get().messages as CoreMessage[],
    text: async ({ content, done }) => {
      // When it's the final content, mark the state as done and ready for the client to access.
      if (done) {
        history.done({
          ...history.get(),
          messages: [
            ...history.get().messages,
            {
              role: "assistant",
              content,
            },
          ],
        });

      return <ChatAssistantMessage content={content} />;
    },
    tools: {
      show_quick_reply_question: {
        description: "Shows the ui for the user to quick reply",
        parameters: z
          .object({
            question: z.string().describe("the question to show"),
            options: z.array(z.string()).describe("the options to show"),
          })
          .required(),
        generate: async function ({ question, options }) {
          console.log(
            "show_quick_reply_question messages",
            history.get().messages
          );

          history.done({
            ...history.get(),
            messages: [
              ...history.get().messages,
              {
                role: "function",
                name: "show_quick_reply_question",
                content: JSON.stringify({ question, options }),
              },
            ],
          });

          return [
            <ChatAssistantMessage content={question} className="mt-2" />,
            <div className="flex flex-col gap-2 mt-2">
              {options.map((option) => (
                <ChatQuickReply key={option} chatQuickReply={option} />
              ))}
            </div>,
          ];
        },
      },
    },
  });
}

Additional context

When the conversations goes through the "text" streaming generation the state seems to be updated sucessfully. But in the case of rendering the questions through the tool 'show_quick_reply_question' (that is what the context says so the assistant uses this tool), after pressing the Button options in the UI a new message is submitted and in that iteration the history that is accessed is not updated and it looks like the first iterations with only 1 initial message and the last ui interface message.

[unstubbable]

The current implementation of streamUI has a known limitation that tool calls overwrite each other and also the text UI (see #1210 for an approach to fix this in the predecessor render). Therefore, you currently need to handle the composition of text and tool UI & state manually, and limit it to a single tool call per message via the system prompt.

Furthermore, make sure that you properly set the assistant content in generate, similar to https://sdk.vercel.ai/docs/ai-sdk-rsc/ai-ui-states#updating-ai-state-on-server, and also the tool content.

Here's a fully working example: https://github.com/unstubbable/mfng-ai-demo/blob/db32ff4c80fc831d62bf07c7603b02e7b9e09891/src/app/submit-user-message.tsx#L86-L119

Disclaimer: This is not a general recommendation of how streamUI is supposed to be used, I'm just sharing the solution that worked for me with the current state of the AI SDK.