feat: initialMessages SSR support (#569)

assistant-ui · Jul 25, 2024 · fc6bc35 · fc6bc35
1 parent 58d5a9a
commit fc6bc35
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 37 deletions.
diff --git a/.changeset/afraid-phones-occur.md b/.changeset/afraid-phones-occur.md
@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react": patch
+---
+
+feat: initialMessages SSR support
diff --git a/packages/react/src/context/providers/ContentPartProvider.tsx b/packages/react/src/context/providers/ContentPartProvider.tsx
@@ -53,9 +53,9 @@ const toContentPartStatus = (
 
 const EMPTY_CONTENT = Object.freeze({ type: "text", text: "" });
 
-const syncContentPart = (
+const getContentPartState = (
   { message }: MessageState,
-  useContentPart: ContentPartContextValue["useContentPart"],
+  useContentPart: ContentPartContextValue["useContentPart"] | undefined,
   partIndex: number,
 ) => {
   let part = message.content[partIndex];
@@ -70,35 +70,44 @@ const syncContentPart = (
 
   // if the content part is the same, don't update
   const status = toContentPartStatus(message, partIndex, part);
-  const currentState = useContentPart.getState();
-  if (currentState.part === part && currentState.status === status) return;
-
-  // sync useContentPart
-  (useContentPart as unknown as StoreApi<ContentPartState>).setState(
-    Object.freeze({
-      part,
-      status,
-    }),
-  );
+  const currentState = useContentPart?.getState();
+  if (
+    currentState &&
+    currentState.part === part &&
+    currentState.status === status
+  )
+    return null;
+
+  return Object.freeze({ part, status });
 };
 
 const useContentPartContext = (partIndex: number) => {
   const { useMessage } = useMessageContext();
   const [context] = useState<ContentPartContextValue>(() => {
     const useContentPart = create<ContentPartState>(
-      () => ({}) as ContentPartState,
+      () => getContentPartState(useMessage.getState(), undefined, partIndex)!,
     );
 
-    syncContentPart(useMessage.getState(), useContentPart, partIndex);
+    getContentPartState(useMessage.getState(), useContentPart, partIndex);
 
     return { useContentPart };
   });
 
   useEffect(() => {
-    syncContentPart(useMessage.getState(), context.useContentPart, partIndex);
-    return useMessage.subscribe((message) => {
-      syncContentPart(message, context.useContentPart, partIndex);
-    });
+    const syncContentPart = (message: MessageState) => {
+      const newState = getContentPartState(
+        message,
+        context.useContentPart,
+        partIndex,
+      );
+      if (!newState) return;
+      (
+        context.useContentPart as unknown as StoreApi<ContentPartState>
+      ).setState(newState, true);
+    };
+
+    syncContentPart(useMessage.getState());
+    return useMessage.subscribe(syncContentPart);
   }, [context, useMessage, partIndex]);
 
   return context;

diff --git a/packages/react/src/context/providers/MessageProvider.tsx b/packages/react/src/context/providers/MessageProvider.tsx
@@ -23,31 +23,30 @@ const getIsLast = (messages: ThreadMessagesState, message: ThreadMessage) => {
   return messages[messages.length - 1]?.id === message.id;
 };
 
-const syncMessage = (
+const getMessageState = (
   messages: ThreadMessagesState,
   getBranches: (messageId: string) => readonly string[],
-  useMessage: MessageContextValue["useMessage"],
+  useMessage: MessageContextValue["useMessage"] | undefined,
   messageIndex: number,
 ) => {
   const parentId = messages[messageIndex - 1]?.id ?? null;
-  const message = messages[messageIndex];
-  if (!message) return;
+  const message = messages[messageIndex]!;
 
   const isLast = getIsLast(messages, message);
   const branches = getBranches(message.id);
 
   // if the message is the same, don't update
-  const currentState = useMessage.getState();
+  const currentState = useMessage?.getState();
   if (
+    currentState &&
     currentState.message === message &&
     currentState.parentId === parentId &&
     currentState.branches === branches &&
     currentState.isLast === isLast
   )
-    return;
+    return null;
 
-  // sync useMessage
-  (useMessage as unknown as StoreApi<MessageState>).setState({
+  return Object.freeze({
     message,
     parentId,
     branches,
@@ -59,7 +58,15 @@ const useMessageContext = (messageIndex: number) => {
   const { useThreadMessages, useThreadActions } = useThreadContext();
 
   const [context] = useState<MessageContextValue>(() => {
-    const useMessage = create<MessageState>(() => ({}) as MessageState);
+    const useMessage = create<MessageState>(
+      () =>
+        getMessageState(
+          useThreadMessages.getState(),
+          useThreadActions.getState().getBranches,
+          undefined,
+          messageIndex,
+        )!,
+    );
     const useMessageUtils = makeMessageUtilsStore();
     const useEditComposer = makeEditComposerStore({
       onEdit: () => {
@@ -92,25 +99,27 @@ const useMessageContext = (messageIndex: number) => {
       },
     });
 
-    syncMessage(
-      useThreadMessages.getState(),
-      useThreadActions.getState().getBranches,
-      useMessage,
-      messageIndex,
-    );
-
     return { useMessage, useMessageUtils, useEditComposer };
   });
 
   useEffect(() => {
-    return useThreadMessages.subscribe((thread) => {
-      syncMessage(
+    const syncMessage = (thread: ThreadMessagesState) => {
+      const newState = getMessageState(
         thread,
         useThreadActions.getState().getBranches,
         context.useMessage,
         messageIndex,
       );
-    });
+      if (!newState) return;
+      (context.useMessage as unknown as StoreApi<MessageState>).setState(
+        newState,
+        true,
+      );
+    };
+
+    syncMessage(useThreadMessages.getState());
+
+    return useThreadMessages.subscribe(syncMessage);
   }, [useThreadMessages, useThreadActions, context, messageIndex]);
 
   return context;