From df6e18bbdf89382e78b75582d33cf1ec8cd82787 Mon Sep 17 00:00:00 2001
From: geoffsee <>
Date: Fri, 30 May 2025 22:27:01 -0400
Subject: [PATCH] Refactor `ClientChatStore` into separate stores for
 modularity and improve maintainability.

---
 src/components/chat/Chat.tsx                  |   4 +-
 src/components/chat/input-menu/InputMenu.tsx  |   2 +-
 src/components/chat/messages/ChatMessages.tsx |   4 +-
 .../chat/messages/MessageBubble.tsx           |   4 +-
 src/stores/ClientChatStore.ts                 | 229 +-----------------
 src/stores/MessagesStore.ts                   |  35 +++
 src/stores/ModelStore.ts                      |  35 +++
 src/stores/StreamStore.ts                     | 121 +++++++++
 src/stores/UIStore.ts                         |  21 ++
 9 files changed, 228 insertions(+), 227 deletions(-)
 create mode 100644 src/stores/MessagesStore.ts
 create mode 100644 src/stores/ModelStore.ts
 create mode 100644 src/stores/StreamStore.ts
 create mode 100644 src/stores/UIStore.ts

diff --git a/src/components/chat/Chat.tsx b/src/components/chat/Chat.tsx
index 9331951..c2e0671 100644
--- a/src/components/chat/Chat.tsx
+++ b/src/components/chat/Chat.tsx
@@ -25,8 +25,8 @@ const Chat = observer(({ height, width }) => {
       width={width}
       gap={0}
     >
-      <GridItem alignSelf="center" hidden={!(chatStore.messages.length < 1)}>
-        <WelcomeHome visible={chatStore.messages.length < 1} />
+      <GridItem alignSelf="center" hidden={!(chatStore.items.length < 1)}>
+        <WelcomeHome visible={chatStore.items.length < 1} />
       </GridItem>
 
       <GridItem
diff --git a/src/components/chat/input-menu/InputMenu.tsx b/src/components/chat/input-menu/InputMenu.tsx
index c10d7f5..8e8e181 100644
--- a/src/components/chat/input-menu/InputMenu.tsx
+++ b/src/components/chat/input-menu/InputMenu.tsx
@@ -71,7 +71,7 @@ const InputMenu: React.FC<{ isDisabled?: boolean }> = observer(
 
     const handleCopyConversation = useCallback(() => {
       navigator.clipboard
-        .writeText(formatConversationMarkdown(ClientChatStore.messages))
+        .writeText(formatConversationMarkdown(ClientchatStore.items))
         .then(() => {
           window.alert(
             "Conversation copied to clipboard. \n\nPaste it somewhere safe!",
diff --git a/src/components/chat/messages/ChatMessages.tsx b/src/components/chat/messages/ChatMessages.tsx
index 1f166ae..b92553c 100644
--- a/src/components/chat/messages/ChatMessages.tsx
+++ b/src/components/chat/messages/ChatMessages.tsx
@@ -27,8 +27,8 @@ const ChatMessages: React.FC<ChatMessagesProps> = observer(({ scrollRef }) => {
         boxShadow="md"
         whiteSpace="pre-wrap"
       >
-        {chatStore.messages.map((msg, index) => {
-          if (index < chatStore.messages.length - 1) {
+        {chatStore.items.map((msg, index) => {
+          if (index < chatStore.items.length - 1) {
             return (
               <GridItem key={index}>
                 <MessageBubble x scrollRef={scrollRef} msg={msg} />
diff --git a/src/components/chat/messages/MessageBubble.tsx b/src/components/chat/messages/MessageBubble.tsx
index 6501bff..50eae3b 100644
--- a/src/components/chat/messages/MessageBubble.tsx
+++ b/src/components/chat/messages/MessageBubble.tsx
@@ -63,12 +63,12 @@ const MessageBubble = observer(({ msg, scrollRef }) => {
 
   useEffect(() => {
     if (
-      clientChatStore.messages.length > 0 &&
+      clientChatStore.items.length > 0 &&
       clientChatStore.isLoading &&
       UserOptionsStore.followModeEnabled
     ) {
       console.log(
-        `${clientChatStore.messages.length}/${clientChatStore.isLoading}/${UserOptionsStore.followModeEnabled}`,
+        `${clientChatStore.items.length}/${clientChatStore.isLoading}/${UserOptionsStore.followModeEnabled}`,
       );
       scrollRef.current?.scrollTo({
         top: scrollRef.current.scrollHeight,
diff --git a/src/stores/ClientChatStore.ts b/src/stores/ClientChatStore.ts
index 84ae381..f34645d 100644
--- a/src/stores/ClientChatStore.ts
+++ b/src/stores/ClientChatStore.ts
@@ -1,229 +1,18 @@
-// ---------------------------
-// stores/MessagesStore.ts
-// ---------------------------
-import { Instance, types } from "mobx-state-tree";
-import Message from "../models/Message";
-
-export const MessagesStore = types
-    .model("MessagesStore", {
-      items: types.optional(types.array(Message), []),
-    })
-    .actions((self) => ({
-      add(message: Instance<typeof Message>) {
-        self.items.push(message);
-      },
-      updateLast(content: string) {
-        if (self.items.length) {
-          self.items[self.items.length - 1].content = content;
-        }
-      },
-      appendLast(content: string) {
-        if (self.items.length) {
-          self.items[self.items.length - 1].content += content;
-        }
-      },
-      removeAfter(index: number) {
-        if (index >= 0 && index < self.items.length) {
-          self.items.splice(index + 1);
-        }
-      },
-      reset() {
-        self.items.clear();
-      },
-    }));
-
-export interface IMessagesStore extends Instance<typeof MessagesStore> {}
-
-// ---------------------------
-// stores/UIStore.ts
-// ---------------------------
-import { Instance, types } from "mobx-state-tree";
-
-export const UIStore = types
-    .model("UIStore", {
-      input: types.optional(types.string, ""),
-      isLoading: types.optional(types.boolean, false),
-    })
-    .actions((self) => ({
-      setInput(value: string) {
-        self.input = value;
-      },
-      setIsLoading(value: boolean) {
-        self.isLoading = value;
-      },
-    }));
-
-export interface IUIStore extends Instance<typeof UIStore> {}
-
-// ---------------------------
-// stores/ModelStore.ts
-// ---------------------------
-import { Instance, types } from "mobx-state-tree";
-
-export const ModelStore = types
-    .model("ModelStore", {
-      model: types.optional(
-          types.string,
-          "meta-llama/llama-4-scout-17b-16e-instruct",
-      ),
-      imageModel: types.optional(types.string, "black-forest-labs/flux-1.1-pro"),
-      supportedModels: types.optional(types.array(types.string), []),
-    })
-    .actions((self) => ({
-      setModel(value: string) {
-        self.model = value;
-        try {
-          localStorage.setItem("recentModel", value);
-        } catch (_) {}
-      },
-      setImageModel(value: string) {
-        self.imageModel = value;
-      },
-      setSupportedModels(list: string[]) {
-        self.supportedModels = list;
-        if (!list.includes(self.model)) {
-          // fall back to last entry (arbitrary but predictable)
-          self.model = list[list.length - 1] ?? self.model;
-        }
-      },
-    }));
-
-export interface IModelStore extends Instance<typeof ModelStore> {}
-
-// ---------------------------
-// stores/StreamStore.ts
-// Handles networking + SSE lifecycle.
-// Depends on MessagesStore, UIStore, and ModelStore via composition.
-// ---------------------------
-import {
-  getParent,
-  Instance,
-  flow,
-  types,
-} from "mobx-state-tree";
-import type { IMessagesStore } from "./MessagesStore";
-import type { IUIStore } from "./UIStore";
-import type { IModelStore } from "./ModelStore";
-import UserOptionsStore from "./UserOptionsStore";
-import Message from "../models/Message";
-
-interface RootDeps extends IMessagesStore, IUIStore, IModelStore {}
-
-export const StreamStore = types
-    .model("StreamStore", {})
-    .volatile(() => ({
-      eventSource: null as EventSource | null,
-    }))
-    .actions((self) => {
-      // helpers
-      const root = getParent<RootDeps>(self);
-
-      function cleanup() {
-        if (self.eventSource) {
-          self.eventSource.close();
-          self.eventSource = null;
-        }
-      }
-
-      const sendMessage = flow(function* () {
-        if (!root.input.trim() || root.isLoading) return;
-        cleanup();
-        yield UserOptionsStore.setFollowModeEnabled(true);
-        root.setIsLoading(true);
-
-        const userMessage = Message.create({
-          content: root.input,
-          role: "user" as const,
-        });
-        root.add(userMessage);
-        root.setInput("");
-
-        try {
-          const payload = { messages: root.items.slice(), model: root.model };
-          // optimistic UI delay (demo‑purpose)
-          yield new Promise((r) => setTimeout(r, 500));
-          root.add(Message.create({ content: "", role: "assistant" }));
-
-          const response: Response = yield fetch("/api/chat", {
-            method: "POST",
-            headers: { "Content-Type": "application/json" },
-            body: JSON.stringify(payload),
-          });
-
-          if (response.status === 429) {
-            root.updateLast("Too many requests • please slow down.");
-            cleanup();
-            return;
-          }
-          if (response.status > 200) {
-            root.updateLast("Error • something went wrong.");
-            cleanup();
-            return;
-          }
-
-          const { streamUrl } = (yield response.json()) as { streamUrl: string };
-          self.eventSource = new EventSource(streamUrl);
-
-          self.eventSource.onmessage = (event) => {
-            try {
-              const parsed = JSON.parse(event.data);
-              if (parsed.type === "error") {
-                root.updateLast(parsed.error);
-                cleanup();
-                root.setIsLoading(false);
-                return;
-              }
-
-              if (
-                  parsed.type === "chat" &&
-                  parsed.data.choices[0]?.finish_reason === "stop"
-              ) {
-                root.appendLast(parsed.data.choices[0]?.delta?.content ?? "");
-                cleanup();
-                root.setIsLoading(false);
-                return;
-              }
-
-              if (parsed.type === "chat") {
-                root.appendLast(parsed.data.choices[0]?.delta?.content ?? "");
-              }
-            } catch (err) {
-              console.error("stream parse error", err);
-            }
-          };
-
-          self.eventSource.onerror = () => cleanup();
-        } catch (err) {
-          console.error("sendMessage", err);
-          root.updateLast("Sorry • network error.");
-          cleanup();
-          root.setIsLoading(false);
-        }
-      });
-
-      const stopIncomingMessage = () => {
-        cleanup();
-        root.setIsLoading(false);
-      };
-
-      return { sendMessage, stopIncomingMessage, cleanup };
-    });
-
-export interface IStreamStore extends Instance<typeof StreamStore> {}
-
 // ---------------------------
 // stores/ClientChatStore.ts (root)
 // ---------------------------
-import { types } from "mobx-state-tree";
-// import { MessagesStore } from "./MessagesStore";
-// import { UIStore } from "./UIStore";
-// import { ModelStore } from "./ModelStore";
-// import { StreamStore } from "./StreamStore";
+import { types, type Instance } from "mobx-state-tree";
+import { MessagesStore } from "./MessagesStore";
+import { UIStore } from "./UIStore";
+import { ModelStore } from "./ModelStore";
+import { StreamStore } from "./StreamStore";
 
 export const ClientChatStore = types
     .compose(MessagesStore, UIStore, ModelStore, StreamStore)
     .named("ClientChatStore");
 
-export const clientChatStore = ClientChatStore.create();
+const clientChatStore = ClientChatStore.create();
 
-export type IClientChatStore = Instance<typeof ClientChatStore>;
\ No newline at end of file
+export type IClientChatStore = Instance<typeof ClientChatStore>;
+
+export default clientChatStore;
diff --git a/src/stores/MessagesStore.ts b/src/stores/MessagesStore.ts
new file mode 100644
index 0000000..b9c73d6
--- /dev/null
+++ b/src/stores/MessagesStore.ts
@@ -0,0 +1,35 @@
+// ---------------------------
+// stores/MessagesStore.ts
+// ---------------------------
+import { Instance, types } from "mobx-state-tree";
+import Message from "../models/Message";
+
+export const MessagesStore = types
+    .model("MessagesStore", {
+        items: types.optional(types.array(Message), []),
+    })
+    .actions((self) => ({
+        add(message: Instance<typeof Message>) {
+            self.items.push(message);
+        },
+        updateLast(content: string) {
+            if (self.items.length) {
+                self.items[self.items.length - 1].content = content;
+            }
+        },
+        appendLast(content: string) {
+            if (self.items.length) {
+                self.items[self.items.length - 1].content += content;
+            }
+        },
+        removeAfter(index: number) {
+            if (index >= 0 && index < self.items.length) {
+                self.items.splice(index + 1);
+            }
+        },
+        reset() {
+            self.items.clear();
+        },
+    }));
+
+export interface IMessagesStore extends Instance<typeof MessagesStore> {}
diff --git a/src/stores/ModelStore.ts b/src/stores/ModelStore.ts
new file mode 100644
index 0000000..07bf79c
--- /dev/null
+++ b/src/stores/ModelStore.ts
@@ -0,0 +1,35 @@
+
+// ---------------------------
+// stores/ModelStore.ts
+// ---------------------------
+import { Instance, types } from "mobx-state-tree";
+
+export const ModelStore = types
+    .model("ModelStore", {
+        model: types.optional(
+            types.string,
+            "meta-llama/llama-4-scout-17b-16e-instruct",
+        ),
+        imageModel: types.optional(types.string, "black-forest-labs/flux-1.1-pro"),
+        supportedModels: types.optional(types.array(types.string), []),
+    })
+    .actions((self) => ({
+        setModel(value: string) {
+            self.model = value;
+            try {
+                localStorage.setItem("recentModel", value);
+            } catch (_) {}
+        },
+        setImageModel(value: string) {
+            self.imageModel = value;
+        },
+        setSupportedModels(list: string[]) {
+            self.supportedModels = list;
+            if (!list.includes(self.model)) {
+                // fall back to last entry (arbitrary but predictable)
+                self.model = list[list.length - 1] ?? self.model;
+            }
+        },
+    }));
+
+export interface IModelStore extends Instance<typeof ModelStore> {}
diff --git a/src/stores/StreamStore.ts b/src/stores/StreamStore.ts
new file mode 100644
index 0000000..fc1a8e9
--- /dev/null
+++ b/src/stores/StreamStore.ts
@@ -0,0 +1,121 @@
+import {
+    getParent,
+    Instance,
+    flow,
+    types,
+} from "mobx-state-tree";
+import type { IMessagesStore } from "./MessagesStore";
+import type { IUIStore } from "./UIStore";
+import type { IModelStore } from "./ModelStore";
+import UserOptionsStore from "./UserOptionsStore";
+import Message from "../models/Message";
+
+interface RootDeps extends IMessagesStore, IUIStore, IModelStore {}
+
+export const StreamStore = types
+    .model("StreamStore", {})
+    .volatile(() => ({
+        eventSource: null as EventSource | null,
+    }))
+    .actions((self: any) => {                  // ← annotate `self` so it isn’t implicitly `any`
+        let root: RootDeps;
+        try {
+            root = getParent<RootDeps>(self);
+        } catch {
+            root = self as any;
+        }
+
+        function cleanup() {
+            if (self.eventSource) {
+                self.eventSource.close();
+                self.eventSource = null;
+            }
+        }
+
+        const sendMessage = flow(function* () {
+            if (!root.input.trim() || root.isLoading) return;
+            cleanup();
+
+            // ← **DO NOT** `yield` a synchronous action
+            UserOptionsStore.setFollowModeEnabled(true);
+            root.setIsLoading(true);
+
+            const userMessage = Message.create({
+                content: root.input,
+                role: "user" as const,
+            });
+            root.add(userMessage);
+            root.setInput("");
+
+            try {
+                const payload = { messages: root.items.slice(), model: root.model };
+
+                yield new Promise((r) => setTimeout(r, 500));
+                root.add(Message.create({ content: "", role: "assistant" }));
+
+                const response: Response = yield fetch("/api/chat", {
+                    method: "POST",
+                    headers: { "Content-Type": "application/json" },
+                    body: JSON.stringify(payload),
+                });
+
+                if (response.status === 429) {
+                    root.updateLast("Too many requests • please slow down.");
+                    cleanup();
+                    return;
+                }
+                if (response.status > 200) {
+                    root.updateLast("Error • something went wrong.");
+                    cleanup();
+                    return;
+                }
+
+                const { streamUrl } = (yield response.json()) as { streamUrl: string };
+                self.eventSource = new EventSource(streamUrl);
+
+                self.eventSource.onmessage = (event: MessageEvent) => {  // ← annotate `event`
+                    try {
+                        const parsed = JSON.parse(event.data);
+                        if (parsed.type === "error") {
+                            root.updateLast(parsed.error);
+                            cleanup();
+                            root.setIsLoading(false);
+                            return;
+                        }
+
+                        if (
+                            parsed.type === "chat" &&
+                            parsed.data.choices[0]?.finish_reason === "stop"
+                        ) {
+                            root.appendLast(parsed.data.choices[0]?.delta?.content ?? "");
+                            cleanup();
+                            root.setIsLoading(false);
+                            return;
+                        }
+
+                        if (parsed.type === "chat") {
+                            root.appendLast(parsed.data.choices[0]?.delta?.content ?? "");
+                        }
+                    } catch (err) {
+                        console.error("stream parse error", err);
+                    }
+                };
+
+                self.eventSource.onerror = () => cleanup();
+            } catch (err) {
+                console.error("sendMessage", err);
+                root.updateLast("Sorry • network error.");
+                cleanup();
+                root.setIsLoading(false);
+            }
+        });
+
+        const stopIncomingMessage = () => {
+            cleanup();
+            root.setIsLoading(false);
+        };
+
+        return { sendMessage, stopIncomingMessage, cleanup };
+    });
+
+export interface IStreamStore extends Instance<typeof StreamStore> {}
\ No newline at end of file
diff --git a/src/stores/UIStore.ts b/src/stores/UIStore.ts
new file mode 100644
index 0000000..e9fc064
--- /dev/null
+++ b/src/stores/UIStore.ts
@@ -0,0 +1,21 @@
+
+// ---------------------------
+// stores/UIStore.ts
+// ---------------------------
+import { Instance, types } from "mobx-state-tree";
+
+export const UIStore = types
+    .model("UIStore", {
+        input: types.optional(types.string, ""),
+        isLoading: types.optional(types.boolean, false),
+    })
+    .actions((self) => ({
+        setInput(value: string) {
+            self.input = value;
+        },
+        setIsLoading(value: boolean) {
+            self.isLoading = value;
+        },
+    }));
+
+export interface IUIStore extends Instance<typeof UIStore> {}