From 7e632b9674a409774c8efa02eed2e9d50fba59e8 Mon Sep 17 00:00:00 2001
From: Ayaan Zaidi <zaidi@uplause.io>
Date: Sat, 21 Feb 2026 17:07:47 +0530
Subject: [PATCH] fix(telegram): split preview per assistant block

---
 src/telegram/bot-message-dispatch.test.ts | 8 ++++----
 src/telegram/bot-message-dispatch.ts      | 3 +++
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/telegram/bot-message-dispatch.test.ts b/src/telegram/bot-message-dispatch.test.ts
index c2f84730938..c58bdc5a81c 100644
--- a/src/telegram/bot-message-dispatch.test.ts
+++ b/src/telegram/bot-message-dispatch.test.ts
@@ -445,7 +445,7 @@ describe("dispatchTelegramMessage draft streaming", () => {
     );
   });
 
-  it("does not force new message for legacy block stream mode", async () => {
+  it("forces new message for next assistant block in legacy block stream mode", async () => {
     const draftStream = createDraftStream(999);
     createTelegramDraftStream.mockReturnValue(draftStream);
     dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
@@ -464,10 +464,10 @@ describe("dispatchTelegramMessage draft streaming", () => {
 
     await dispatchWithContext({ context: createContext(), streamMode: "block" });
 
-    expect(draftStream.forceNewMessage).not.toHaveBeenCalled();
+    expect(draftStream.forceNewMessage).toHaveBeenCalledTimes(1);
   });
 
-  it("does not force new message in partial mode when assistant message restarts", async () => {
+  it("forces new message in partial mode when assistant message restarts", async () => {
     const draftStream = createDraftStream(999);
     createTelegramDraftStream.mockReturnValue(draftStream);
     dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
@@ -483,7 +483,7 @@ describe("dispatchTelegramMessage draft streaming", () => {
 
     await dispatchWithContext({ context: createContext(), streamMode: "partial" });
 
-    expect(draftStream.forceNewMessage).not.toHaveBeenCalled();
+    expect(draftStream.forceNewMessage).toHaveBeenCalledTimes(1);
   });
 
   it("does not force new message on first assistant message start", async () => {
diff --git a/src/telegram/bot-message-dispatch.ts b/src/telegram/bot-message-dispatch.ts
index 7026b8cca6c..ed02d700d62 100644
--- a/src/telegram/bot-message-dispatch.ts
+++ b/src/telegram/bot-message-dispatch.ts
@@ -630,6 +630,9 @@ export const dispatchTelegramMessage = async ({
         onAssistantMessageStart: answerLane.stream
           ? () => {
               reasoningStepState.resetForNextStep();
+              if (answerLane.hasStreamedMessage) {
+                answerLane.stream?.forceNewMessage();
+              }
               resetDraftLaneState(answerLane);
             }
           : undefined,