How to load messages on client side using shared chat context?

[hoon]

here’s my code:
the problem is when i set chat this way, the new message that child component get won’t be updated except the first time and the finish time

import { useParams } from "@umijs/max";
import { type CustomTransport } from "@/libs/ai/transport";
import { getChatHistory } from "@/libs/api";
import { Chat } from "@ai-sdk/react";
import { useQuery } from "@tanstack/react-query";
import type { UIMessage } from "ai";
import {
	type ReactNode,
	createContext,
	useContext,
	useEffect,
	useState,
} from "react";

interface ChatContextValue {
	chat: Chat<UIMessage>;
	clearChat: () => void;
	transport: CustomTransport;

	loadingMessage?: boolean;
}

const ChatContext = createContext<ChatContextValue | undefined>(undefined);

const createChat = (opts: ChatOptsType) => new Chat<UIMessage>(opts);
type ChatOptsType = ConstructorParameters<typeof Chat>[0];
export const ChatProvider = ({
	children,
	chatOpts,
}: {
	children: ReactNode;
	chatOpts: ChatOptsType & {
		transport: CustomTransport;
	};
}) => {
	const { id } = useParams();
	const [chat, setChat] = useState(() => createChat(chatOpts));

	const { data, isSuccess } = useQuery({
		queryKey: ["chatHistory", id],
		queryFn: () => getChatHistory(id!),
		retry: false,
	});

	useEffect(() => {
		if (isSuccess ) {
			setChat(
				createChat({
					...chatOpts,
					messages: data.messages,
				}),
			);
		}
	}, [isSuccess, data, chatOpts, setChat]);

	const clearChat = () => {
		setChat(createChat(chatOpts));
	};

	return (
		<ChatContext.Provider
			value={{
				chat,
				clearChat,
				transport: chatOpts.transport,
			}}
		>
			{children}
		</ChatContext.Provider>
	);
};

export function useSharedChatContext() {
	const context = useContext(ChatContext);
	if (!context) {
		throw new Error("useSharedChatContext must be used within a ChatProvider");
	}
	return context;
}

[amyegan]

I asked v0 to review the code and identify the problem. Here’s what it gave me:

Looking at your code, the issue is that you’re recreating the entire Chat instance every time the query succeeds, which breaks the reactivity of the chat state. When you create a new Chat instance, you lose all the internal state management and subscriptions that child components depend on.

Here’s how you can fix this:

Option 1: Initialize messages once and let Chat manage updates

export const ChatProvider = ({
  children,
  chatOpts,
}: {
  children: ReactNode;
  chatOpts: ChatOptsType & {
    transport: CustomTransport;
  };
}) => {
  const { id } = useParams();
  const [chat, setChat] = useState<Chat<UIMessage> | null>(null);

  const { data, isSuccess } = useQuery({
    queryKey: ["chatHistory", id],
    queryFn: () => getChatHistory(id!),
    retry: false,
  });

  useEffect(() => {
    if (isSuccess && data) {
      // Only create chat once when data is loaded
      const newChat = createChat({
        ...chatOpts,
        messages: data.messages,
      });
      setChat(newChat);
    }
  }, [isSuccess, data?.messages]); // Remove chatOpts from dependencies

  const clearChat = () => {
    setChat(createChat(chatOpts));
  };

  if (!chat) {
    return <div>Loading...</div>; // or your loading component
  }

  return (
    <ChatContext.Provider
      value={{
        chat,
        clearChat,
        transport: chatOpts.transport,
      }}
    >
      {children}
    </ChatContext.Provider>
  );
};

Option 2: Use the Chat’s built-in methods to update messages

useEffect(() => {
  if (isSuccess && data && chat) {
    // Use chat's setMessages method instead of recreating
    chat.setMessages(data.messages);
  }
}, [isSuccess, data, chat]);

The key issue was that chatOpts in your dependency array was causing the effect to run on every render, recreating the chat instance and breaking the connection with child components.

I hope that helps!

[amyegan]

Hey @hoon! Just checking in to see if you still need help with loading messages in your shared chat context. Have you found a solution, or do you want to dive deeper into this together? Let’s figure this out!