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

(edited)

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;
}

Vercel Staff

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!

Vercel Staff

Hey ! 😊 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!

This topic is closed. It no longer accepts new replies.

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; }

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> ); };