(EAI-629) [UI] Switch state management to Zustand (#581)

Author: nlarew

package-lock.json

@@ -42,7 +42,7 @@
   },
   "scripts": {
     "build:dts": "tsc --project tsconfig.declarations.json",
-    "build": "npm run build:component && npm run build:component:cjs && if [ \"$ENVIRONMENT\" = \"qa\" ]; then npm run build:static-site:qa && echo \"qa\"; else npm run build:static-site:staging; fi",
+    "build": "npm run build:component && if [ \"$ENVIRONMENT\" = \"qa\" ]; then npm run build:static-site:qa && echo \"qa\"; else npm run build:static-site:staging; fi",
     "build:component": "npm run build:component:esm && npm run build:component:cjs && npm run build:dts",
     "build:component:esm": "tsc --project tsconfig.build.json && VITE_MODULE_FORMAT=es vite build --config vite.config.component.js",
     "build:component:cjs": "tsc --project tsconfig.build.json && VITE_MODULE_FORMAT=cjs vite build --config vite.config.component.js",
@@ -107,7 +107,8 @@
     "remark": "^15.0.1",
     "remark-parse": "^11.0.0",
     "unified": "^11.0.4",
-    "unist-util-visit": "^5.0.0"
+    "unist-util-visit": "^5.0.0",
+    "zustand": "^5.0.2"
   },
   "devDependencies": {
     "@rollup/plugin-inject": "^5.0.3",

packages/mongodb-chatbot-ui/package.json

@@ -59,10 +59,10 @@ export function ChatMessageFeed(props: ChatMessageFeedProps) {
         </DisclaimerText>
       ) : null}
       {messages.map((message, idx) => {
-        const isLoading = conversation.isStreamingMessage
-          ? message.id === conversation.streamingMessage?.id &&
-            conversation.streamingMessage?.content === ""
-          : false;
+        const isLoading =
+          message.id === conversation.streamingMessageId &&
+          conversation.getMessage(conversation.streamingMessageId ?? "")
+            ?.content === "";
 
         const isInitialMessage = idx === 0;
 
@@ -76,8 +76,7 @@ export function ChatMessageFeed(props: ChatMessageFeedProps) {
               message.role === "assistant" &&
               !isLoading &&
               !(
-                awaitingReply &&
-                conversation.streamingMessage?.id === message.id
+                awaitingReply && conversation.streamingMessageId === message.id
               ) &&
               // We don't want users to rate the initial message (and they can't because it's not in the database)
               !isInitialMessage

packages/mongodb-chatbot-ui/src/ChatMessageFeed.tsx

@@ -6,7 +6,7 @@ import { useChatbot, OpenCloseHandlers, UseChatbotProps } from "./useChatbot";
 import { LinkDataProvider } from "./LinkDataProvider";
 import { type User } from "./useUser";
 import { ChatbotProvider } from "./ChatbotProvider";
-import ConversationProvider from "./ConversationProvider";
+import { ConversationStateProvider } from "./ConversationStateProvider";
 import { RenameFields } from "./utils";
 import { HotkeyContextProvider } from "./HotkeyContext";
 
@@ -41,34 +41,57 @@ export function Chatbot({
   const maxCommentCharacters =
     props.maxCommentCharacters ?? DEFAULT_MAX_COMMENT_CHARACTERS;
 
-  const chatbotData = useChatbot({
-    chatbotName: name,
-    serverBaseUrl,
-    shouldStream,
-    fetchOptions,
-    isExperimental,
-    maxInputCharacters,
-    maxCommentCharacters,
-    onOpen,
-    onClose,
-    sortMessageReferences,
-  });
-
   const tck = props.tck ?? "mongodb_ai_chatbot";
 
   return (
     <LeafyGreenProvider darkMode={darkMode}>
-      <LinkDataProvider tck={tck}>
-        <UserProvider user={user}>
-          <ChatbotProvider {...chatbotData}>
-            <HotkeyContextProvider>
-              <ConversationProvider conversation={chatbotData.conversation}>
-                {children}
-              </ConversationProvider>
-            </HotkeyContextProvider>
-          </ChatbotProvider>
-        </UserProvider>
-      </LinkDataProvider>
+      <ConversationStateProvider>
+        <LinkDataProvider tck={tck}>
+          <UserProvider user={user}>
+            <InnerChatbot
+              fetchOptions={fetchOptions}
+              isExperimental={isExperimental}
+              maxCommentCharacters={maxCommentCharacters}
+              maxInputCharacters={maxInputCharacters}
+              name={name}
+              onOpen={onOpen}
+              onClose={onClose}
+              serverBaseUrl={serverBaseUrl}
+              shouldStream={shouldStream}
+              sortMessageReferences={sortMessageReferences}
+            >
+              {children}
+            </InnerChatbot>
+          </UserProvider>
+        </LinkDataProvider>
+      </ConversationStateProvider>
     </LeafyGreenProvider>
   );
 }
+
+type InnerChatbotProps = Pick<
+  ChatbotProps,
+  | "children"
+  | "fetchOptions"
+  | "isExperimental"
+  | "maxCommentCharacters"
+  | "maxInputCharacters"
+  | "name"
+  | "onOpen"
+  | "onClose"
+  | "serverBaseUrl"
+  | "shouldStream"
+  | "sortMessageReferences"
+>;
+
+function InnerChatbot({ children, ...props }: InnerChatbotProps) {
+  const chatbotData = useChatbot({
+    ...props,
+  });
+
+  return (
+    <ChatbotProvider {...chatbotData}>
+      <HotkeyContextProvider>{children}</HotkeyContextProvider>
+    </ChatbotProvider>
+  );
+}

packages/mongodb-chatbot-ui/src/Chatbot.tsx

@@ -9,10 +9,10 @@ export type ChatbotProviderProps = ChatbotData & {
 
 export function ChatbotProvider({
   children,
-  ...linkData
+  ...chatbotData
 }: ChatbotProviderProps) {
   return (
-    <ChatbotContext.Provider value={linkData}>
+    <ChatbotContext.Provider value={chatbotData}>
       {children}
     </ChatbotContext.Provider>
   );

packages/mongodb-chatbot-ui/src/ChatbotProvider.tsx

@@ -0,0 +1,48 @@
+import { createContext, useState } from "react";
+import {
+  ConversationStore,
+  initialConversationState,
+  makeConversationStore,
+} from "./conversationStore";
+
+const ConversationStoreNotInitialized =
+  "A conversation store was used before it was initialized.";
+
+const defaultState = {
+  ...initialConversationState,
+  api: new Proxy(
+    {} as ReturnType<ConversationStore["getInitialState"]>["api"],
+    {
+      get: () => () => {
+        throw new Error(ConversationStoreNotInitialized);
+      },
+    }
+  ),
+} as ReturnType<ConversationStore["getInitialState"]>;
+
+const defaultStore = {
+  getInitialState: () => defaultState,
+  getState: () => defaultState,
+  setState: () => {
+    throw new Error(ConversationStoreNotInitialized);
+  },
+  subscribe: () => () => {
+    throw new Error(ConversationStoreNotInitialized);
+  },
+} satisfies ConversationStore;
+
+export const ConversationStateContext =
+  createContext<ConversationStore>(defaultStore);
+
+export function ConversationStateProvider(props: {
+  children: React.ReactNode;
+}) {
+  // The useState initializer function ensures that the store is only created once per mount of this provider.
+  // See https://tkdodo.eu/blog/use-state-for-one-time-initializations for more information.
+  const [store] = useState(() => makeConversationStore("state"));
+  return (
+    <ConversationStateContext.Provider value={store}>
+      {props.children}
+    </ConversationStateContext.Provider>
+  );
+}

packages/mongodb-chatbot-ui/src/ConversationProvider.tsx

@@ -4,11 +4,12 @@ import {
   FloatingActionButtonTrigger,
   type FloatingActionButtonTriggerProps,
 } from "./FloatingActionButtonTrigger";
-import { ModalView, ModalViewProps } from "./ModalView";
+import { type ModalViewProps } from "./ModalView";
 import { MongoDbLegalDisclosure } from "./MongoDbLegal";
 import { mongoDbVerifyInformationMessage } from "./ui-text";
 import { PoweredByAtlasVectorSearch } from "./PoweredByAtlasVectorSearch";
 import { css } from "@emotion/css";
+import { lazy, Suspense } from "react";
 import { References } from "./references";
 
 export type DevCenterChatbotProps = DarkModeProps & {
@@ -17,6 +18,10 @@ export type DevCenterChatbotProps = DarkModeProps & {
   initialMessageReferences?: References;
 };
 
+const ModalView = lazy(() =>
+  import("./ModalView").then((m) => ({ default: m.ModalView }))
+);
+
 export function DevCenterChatbot(props: DevCenterChatbotProps) {
   const { darkMode } = useDarkMode(props.darkMode);
 
@@ -46,7 +51,9 @@ export function DevCenterChatbot(props: DevCenterChatbotProps) {
   return (
     <>
       <FloatingActionButtonTrigger {...triggerProps} />
-      <ModalView {...viewProps} />
+      <Suspense fallback={null}>
+        <ModalView {...viewProps} />
+      </Suspense>
     </>
   );
 }

packages/mongodb-chatbot-ui/src/ConversationStateProvider.tsx

@@ -1,9 +1,14 @@
 import { useDarkMode } from "@leafygreen-ui/leafygreen-provider";
 import { DarkModeProps } from "./DarkMode";
 import { InputBarTrigger, InputBarTriggerProps } from "./InputBarTrigger";
-import { ModalView, ModalViewProps } from "./ModalView";
+import { type ModalViewProps } from "./ModalView";
 import { MongoDbLegalDisclosure } from "./MongoDbLegal";
 import { mongoDbVerifyInformationMessage } from "./ui-text";
+import { lazy } from "react";
+
+const ModalView = lazy(() =>
+  import("./ModalView").then((m) => ({ default: m.ModalView }))
+);
 
 export type DocsChatbotProps = DarkModeProps & {
   suggestedPrompts?: string[];

packages/mongodb-chatbot-ui/src/DevCenterChatbot.tsx

@@ -65,6 +65,8 @@ export function ModalView(props: ModalViewProps) {
 
   const chatWindowInputBarId = inputBarId ?? "chatbot-modal-input-bar";
 
+  console.log("chatWindowInputBarId", chatWindowInputBarId);
+
   return (
     <Suspense fallback={null}>
       {open ? (
@@ -77,12 +79,7 @@ export function ModalView(props: ModalViewProps) {
           }
           shouldClose={shouldClose}
         >
-          <Suspense fallback={null}>
-            <ChatWindow
-              inputBarId={chatWindowInputBarId}
-              {...chatWindowProps}
-            />
-          </Suspense>
+          <ChatWindow inputBarId={chatWindowInputBarId} {...chatWindowProps} />
         </Modal>
       ) : null}
     </Suspense>