From e3568f9a973848519c4470a345f85c3fd73aa224 Mon Sep 17 00:00:00 2001
From: Will Chen <willchen90@gmail.com>
Date: Fri, 18 Apr 2025 14:00:05 -0700
Subject: [PATCH] update message approval status after processing message

---
 src/ipc/handlers/chat_stream_handlers.ts | 21 +++++++++++++++------
 src/ipc/handlers/proposal_handlers.ts    | 12 ++++--------
 src/ipc/processors/response_processor.ts | 16 ++++++++++++++--
 3 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/src/ipc/handlers/chat_stream_handlers.ts b/src/ipc/handlers/chat_stream_handlers.ts
index adb983f..3fcb7d6 100644
--- a/src/ipc/handlers/chat_stream_handlers.ts
+++ b/src/ipc/handlers/chat_stream_handlers.ts
@@ -261,18 +261,27 @@ export function registerChatStreamHandlers() {
           })
           .returning();
 
-        await db
-          .update(messages)
-          .set({ content: fullResponse })
-          .where(eq(messages.id, assistantMessage.id));
-
         if (readSettings().autoApproveChanges) {
           const status = await processFullResponseActions(
             fullResponse,
             req.chatId,
-            { chatSummary }
+            { chatSummary, messageId: assistantMessage.id }
           );
 
+          const chat = await db.query.chats.findFirst({
+            where: eq(chats.id, req.chatId),
+            with: {
+              messages: {
+                orderBy: (messages, { asc }) => [asc(messages.createdAt)],
+              },
+            },
+          });
+
+          event.sender.send("chat:response:chunk", {
+            chatId: req.chatId,
+            messages: chat!.messages,
+          });
+
           if (status.error) {
             event.sender.send(
               "chat:response:error",
diff --git a/src/ipc/handlers/proposal_handlers.ts b/src/ipc/handlers/proposal_handlers.ts
index d5ef626..08bac4a 100644
--- a/src/ipc/handlers/proposal_handlers.ts
+++ b/src/ipc/handlers/proposal_handlers.ts
@@ -130,7 +130,10 @@ const approveProposalHandler = async (
     const processResult = await processFullResponseActions(
       messageToApprove.content,
       chatId,
-      { chatSummary: chatSummary ?? undefined } // Pass summary if found
+      {
+        chatSummary: chatSummary ?? undefined,
+        messageId,
+      } // Pass summary if found
     );
 
     if (processResult.error) {
@@ -146,13 +149,6 @@ const approveProposalHandler = async (
       };
     }
 
-    // 3. Update the message's approval state to 'approved'
-    await db
-      .update(messages)
-      .set({ approvalState: "approved" })
-      .where(eq(messages.id, messageId));
-
-    console.log(`Message ${messageId} marked as approved.`);
     return { success: true };
   } catch (error) {
     console.error(
diff --git a/src/ipc/processors/response_processor.ts b/src/ipc/processors/response_processor.ts
index b18cdb6..8eea20d 100644
--- a/src/ipc/processors/response_processor.ts
+++ b/src/ipc/processors/response_processor.ts
@@ -1,5 +1,5 @@
 import { db } from "../../db";
-import { chats } from "../../db/schema";
+import { chats, messages } from "../../db/schema";
 import { eq } from "drizzle-orm";
 import fs from "node:fs";
 import { getDyadAppPath } from "../../paths/paths";
@@ -100,7 +100,10 @@ export function getDyadChatSummaryTag(fullResponse: string): string | null {
 export async function processFullResponseActions(
   fullResponse: string,
   chatId: number,
-  { chatSummary }: { chatSummary: string | undefined }
+  {
+    chatSummary,
+    messageId,
+  }: { chatSummary: string | undefined; messageId: number }
 ): Promise<{ updatedFiles?: boolean; error?: string }> {
   // Get the app associated with the chat
   const chatWithApp = await db.query.chats.findFirst({
@@ -248,6 +251,15 @@ export async function processFullResponseActions(
         author: await getGitAuthor(),
       });
       console.log(`Successfully committed changes: ${changes.join(", ")}`);
+
+      // Update the message to approved
+      await db
+        .update(messages)
+        .set({
+          approvalState: "approved",
+        })
+        .where(eq(messages.id, messageId));
+
       return { updatedFiles: true };
     }