From 74e8aa2d99e6a4efc8f62f6fdb847f4393897a1c Mon Sep 17 00:00:00 2001
From: ranade-oss <ranade.abhijit@gmail.com>
Date: Fri, 17 Jul 2026 02:20:09 -0400
Subject: [PATCH]  Finalize live assistant responses from saved stream

---
 backend/src/routes/chat.ts                    | 15 ++++++++++
 backend/src/routes/projectChat.ts             | 12 ++++++++
 frontend/src/app/hooks/useAssistantChat.ts    | 30 +++++++++++++++++++
 tests/baseline/ross-chat-persistence.test.mjs | 18 +++++++++++
 4 files changed, 75 insertions(+)

diff --git a/backend/src/routes/chat.ts b/backend/src/routes/chat.ts
index 17b612092..5a50023b4 100644
--- a/backend/src/routes/chat.ts
+++ b/backend/src/routes/chat.ts
@@ -730,6 +730,21 @@ chatRouter.post("/", requireAuth, async (req, res) => {
             });
         }
 
+        // Send one canonical post-persistence snapshot. The incremental
+        // stream deliberately holds back a short tail while checking for the
+        // <CITATIONS> marker. Some proxies/browsers can render the earlier
+        // chunks but miss that final small delta. Replacing the provisional
+        // client events with this saved snapshot guarantees that the live UI
+        // and a refreshed chat show the same complete answer.
+        write(
+            `data: ${JSON.stringify({
+                type: "assistant_message_final",
+                events: persistedEvents,
+                citations,
+            })}\n\n`,
+        );
+        write("data: [DONE]\n\n");
+
         if (!chatTitle && lastUser?.content) {
             await db
                 .from("chats")
diff --git a/backend/src/routes/projectChat.ts b/backend/src/routes/projectChat.ts
index 4d88a1e78..d4d6d20bd 100644
--- a/backend/src/routes/projectChat.ts
+++ b/backend/src/routes/projectChat.ts
@@ -309,6 +309,18 @@ projectChatRouter.post("/", requireAuth, async (req, res) => {
             });
         }
 
+        // Send the canonical saved response after persistence so the browser
+        // can replace its provisional stream, including any short tail held
+        // back while the server checked for the <CITATIONS> marker.
+        write(
+            `data: ${JSON.stringify({
+                type: "assistant_message_final",
+                events: persistedEvents,
+                citations,
+            })}\n\n`,
+        );
+        write("data: [DONE]\n\n");
+
         if (!chatTitle && lastUser?.content) {
             await db
                 .from("chats")
diff --git a/frontend/src/app/hooks/useAssistantChat.ts b/frontend/src/app/hooks/useAssistantChat.ts
index c3eb965d7..cc42a3e82 100644
--- a/frontend/src/app/hooks/useAssistantChat.ts
+++ b/frontend/src/app/hooks/useAssistantChat.ts
@@ -413,6 +413,36 @@ export function useAssistantChat({
               continue;
             }
 
+            if (data.type === "assistant_message_final") {
+              const finalEvents = Array.isArray(data.events)
+                ? (data.events as AssistantEvent[])
+                : [];
+              const finalCitations = Array.isArray(data.citations)
+                ? (data.citations as Citation[])
+                : [];
+
+              eventsRef.current = finalEvents;
+              updateLatestAssistantMessage((assistantMessage) => ({
+                ...assistantMessage,
+                content: finalEvents
+                  .filter(
+                    (
+                      event,
+                    ): event is Extract<
+                      AssistantEvent,
+                      { type: "content" }
+                    > => event.type === "content",
+                  )
+                  .map((event) => event.text)
+                  .join(""),
+                events: finalEvents,
+                citations: finalCitations,
+                citationStatus: finalCitations.length ? "final" : undefined,
+              }));
+              setIsLoadingCitations(false);
+              continue;
+            }
+
             if (data.type === "error") {
               const message = readableStreamError(data.message);
               clearStreamingPlaceholders();
diff --git a/tests/baseline/ross-chat-persistence.test.mjs b/tests/baseline/ross-chat-persistence.test.mjs
index 8eeb20344..0a629272a 100644
--- a/tests/baseline/ross-chat-persistence.test.mjs
+++ b/tests/baseline/ross-chat-persistence.test.mjs
@@ -43,3 +43,21 @@ test("the live assistant view reconciles with persisted messages", () => {
   assert.match(hook, /setMessages\(persistedMessages\)/);
   assert.match(hook, /finalizeStreamingContent\(\);[\s\S]*?getChat\(finalChatId\)/);
 });
+
+test("chat streams finish with a canonical post-persistence snapshot", () => {
+  const routes = [
+    read("backend/src/routes/chat.ts"),
+    read("backend/src/routes/projectChat.ts"),
+  ];
+  const hook = read("frontend/src/app/hooks/useAssistantChat.ts");
+
+  for (const route of routes) {
+    assert.match(
+      route,
+      /persistedEvents[\s\S]*?type: "assistant_message_final"[\s\S]*?events: persistedEvents/,
+    );
+  }
+  assert.match(hook, /data\.type === "assistant_message_final"/);
+  assert.match(hook, /eventsRef\.current = finalEvents/);
+  assert.match(hook, /events: finalEvents/);
+});