fix: type error

2 files changed, 38 insertions, 9 deletions

diff --git a/packages/app/src/components/prompt-input.tsx b/packages/app/src/components/prompt-input.tsx
index e4a5a5150..5c64b87cc 100644
--- a/packages/app/src/components/prompt-input.tsx
+++ b/packages/app/src/components/prompt-input.tsx
@@ -1388,10 +1388,27 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
       model,
     }
 
-    const setSyncStore = sessionDirectory === projectDirectory ? sync.set : globalSync.child(sessionDirectory)[1]
-
     const addOptimisticMessage = () => {
-      setSyncStore(
+      if (sessionDirectory === projectDirectory) {
+        sync.set(
+          produce((draft) => {
+            const messages = draft.message[session.id]
+            if (!messages) {
+              draft.message[session.id] = [optimisticMessage]
+            } else {
+              const result = Binary.search(messages, messageID, (m) => m.id)
+              messages.splice(result.index, 0, optimisticMessage)
+            }
+            draft.part[messageID] = optimisticParts
+              .filter((p) => !!p?.id)
+              .slice()
+              .sort((a, b) => a.id.localeCompare(b.id))
+          }),
+        )
+        return
+      }
+
+      globalSync.child(sessionDirectory)[1](
         produce((draft) => {
           const messages = draft.message[session.id]
           if (!messages) {
@@ -1409,7 +1426,21 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     }
 
     const removeOptimisticMessage = () => {
-      setSyncStore(
+      if (sessionDirectory === projectDirectory) {
+        sync.set(
+          produce((draft) => {
+            const messages = draft.message[session.id]
+            if (messages) {
+              const result = Binary.search(messages, messageID, (m) => m.id)
+              if (result.found) messages.splice(result.index, 1)
+            }
+            delete draft.part[messageID]
+          }),
+        )
+        return
+      }
+
+      globalSync.child(sessionDirectory)[1](
         produce((draft) => {
           const messages = draft.message[session.id]
           if (messages) {
diff --git a/packages/app/src/context/sync.tsx b/packages/app/src/context/sync.tsx
index 96872584e..58744045b 100644
--- a/packages/app/src/context/sync.tsx
+++ b/packages/app/src/context/sync.tsx
@@ -101,15 +101,13 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
         })
     }
 
-    const set: (...args: Parameters<Setter>) => ReturnType<Setter> = (...args) => {
-      return current()[1](...args)
-    }
-
     return {
       get data() {
         return current()[0]
       },
-      set,
+      get set(): Setter {
+        return current()[1]
+      },
       get status() {
         return current()[0].status
       },