onGetUIState does not keep AIState in sync with UIState when searchParams change

[james-pratama]

Hi there! I am having issues with keeping my UIState stay up to date with the latest changes to AIState. Specifically, onGetUIState does not seem to re-render UIState with the latest AIState when AIState has been router.refreshed() with new data due to changes in the searchParam.

In my code, I have a list of documents, each of which can have many chats. Currently, when I switch between documents (dynamic route), the AIState and UIState properly updates and I can see the most recent chat and its messages for that document.

However, when I switch between chats in the same document, which only changes the chat searchParams and triggers server component to fetch new messages for the clicked chat, I can see that the new messages are correctly getting passed into initialAIState to update AIState in the onGetUIState. However, the change in AIState does not immediately refresh the UIState. The UIState stays stale until I refresh the page or navigate to another document. The UIState does not re-render if only the searchParams changed, even though I added router.refresh() and force-dynamic.

Below you can see what happens when I switch between chats in the same document in AIState and UIState. My AIState first has a length of 2 messages, and then changes 10 after I click on a new chat in the list of the document's chats.

However, the length of messages in the UIState stays the same at 2 even after I've clicked on the new chat. It only updates when I click on a completely new document.

The root of the problem seems to be that here inside onGetUIState, even when AIState changes because initialAIState is updated in AI Provider, this does not seem to properly update UIState even though its returned with the new data. It only updates when I render a new document or refresh the page.

Anyone know if I'm missing something obvious here?

Url looks like this:
http://localhost:3000/draft/1ec40e4a-edf1-4757-8bfb-9db398fa1fd4?chat=4f956cb0-7e03-4015-8231-9fc01f324fe6

  onGetUIState: async () => {
    "use server";

    const aiState = getAIState();

    return getUIStateFromAIState(aiState);
  },

Console.log of AIState in action.tsx

Console.log of UIState in client chat messages

My handler when to update searchParams when I switch chats in a document

const updateChatParam = useCallback(
  async (chatId: string | null) => {
    const params = new URLSearchParams(searchParams.toString());

    if (chatId === null) {
      params.set("chat", "old");
    } else {
      params.set("chat", chatId);
    }
    router.push(
      pathname + (params.toString() ? "?" + params.toString() : ""),
    );
    router.refresh();
  },
  [router, params_draft_id, searchParams, pathname],
);

My server component where I wrap my chat around AI

import { createClient } from "@/lib/supabase/server";
import TextEditor from "@/app/(authorized)/components/textEditor";
import Chat from "@/app/(authorized)/components/chat";
import { AI } from "@/app/actions";
import { revalidatePath } from "next/cache";

export const runtime = "nodejs";
export const maxDuration = 300;
export const dynamic = "force-dynamic";
export const revalidate = 0;

export default async function Page({ params, searchParams }) {
  const supabase = createClient();
  const {
    data: { user },
  } = await supabase.auth.getUser();
  const draft_id = params.draft_id;

  async function findChatId() {
    if (searchParams.chat) {
      return searchParams.chat;
    } else {
      const { chatId } = await getChat(supabase, user, draft_id);
      return chatId;
    }
  }

  const chat_id = await findChatId();

  let initial_messages = [];

  if (chat_id !== "old") {
    const { initialMessages } = await getMessages(supabase, user, chat_id);

    initial_messages = initialMessages;
  }
  return (
    <div className="h-full w-full flex flex-row relative">
      <AI initialAIState={{ chat_id: chat_id, messages: initial_messages }}>
        <TextEditor user={user} />
        <Chat user={user} />
      </AI>
    </div>
  );
}

My onGetUIState in action.tsx

export const AI = createAI<AIState, UIState>({
  actions: {
    getResponse,
  },
  onGetUIState: async () => {
    "use server";

    const aiState = getAIState();

    return getUIStateFromAIState(aiState);
  },
  onSetAIState: async ({ state, done }) => {
    "use server";
  },
});

export const getUIStateFromAIState = (aiState: AIState) => {
  return aiState.messages
    .map((message) => {
      const {
        role,
        content,
        id,
        created_at,
        message_group_id,
        response_type,
      } = message;
      switch (role) {
        case "user":
          return {
            message_id: id,
            created_at,
            message_group_id,
            component: <UserMessage message={JSON.stringify(content)} />,
          };
        case "assistant":
          const answer = createStreamableValue();
          answer.done(content);
          switch (response_type) {
            case "WriteAnswer":
              return {
                message_id: id,
                created_at,
                message_group_id,
                component: (
                  <>
                    <AIResponseHeader />
                    <AIWriteResponse content={answer.value} />
                  </>
                ),
              };
          }
        default:
          return {
            message_id: id,
            created_at,
            message_group_id,
            component: null,
          };
      }
    })
    .filter((message) => message !== null) as UIState;
};

[james-pratama]

Hey @lgrammel, I am a huge huge fan of your work 🙏 Any chance you have some quick tips on how to make AIState update UIState when server component and initialAIState is refreshed by new SearchParams?