From 1a1d0088ad9ebd2c443c66e4e77c35c87fbd1b4a Mon Sep 17 00:00:00 2001
From: Ayaan Zaidi <zaidi@uplause.io>
Date: Tue, 17 Feb 2026 12:04:22 +0530
Subject: [PATCH] fix(telegram): keep partial stream previews in one message

---
 src/telegram/bot-message-dispatch.test.ts | 38 +++++++++++++++++++++++
 src/telegram/bot-message-dispatch.ts      | 18 +++++------
 2 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/src/telegram/bot-message-dispatch.test.ts b/src/telegram/bot-message-dispatch.test.ts
index 3c02a1f775f..8893628fd17 100644
--- a/src/telegram/bot-message-dispatch.test.ts
+++ b/src/telegram/bot-message-dispatch.test.ts
@@ -343,6 +343,25 @@ describe("dispatchTelegramMessage draft streaming", () => {
     expect(draftStream.forceNewMessage).toHaveBeenCalled();
   });
 
+  it("does not force new message in partial mode when assistant message restarts", async () => {
+    const draftStream = createDraftStream(999);
+    createTelegramDraftStream.mockReturnValue(draftStream);
+    dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
+      async ({ dispatcherOptions, replyOptions }) => {
+        await replyOptions?.onPartialReply?.({ text: "First response" });
+        await replyOptions?.onAssistantMessageStart?.();
+        await replyOptions?.onPartialReply?.({ text: "After tool call" });
+        await dispatcherOptions.deliver({ text: "After tool call" }, { kind: "final" });
+        return { queuedFinal: true };
+      },
+    );
+    deliverReplies.mockResolvedValue({ delivered: true });
+
+    await dispatchWithContext({ context: createContext(), streamMode: "partial" });
+
+    expect(draftStream.forceNewMessage).not.toHaveBeenCalled();
+  });
+
   it("does not force new message on first assistant message start", async () => {
     const draftStream = createDraftStream(999);
     createTelegramDraftStream.mockReturnValue(draftStream);
@@ -390,6 +409,25 @@ describe("dispatchTelegramMessage draft streaming", () => {
     expect(draftStream.forceNewMessage).toHaveBeenCalled();
   });
 
+  it("does not force new message in partial mode when reasoning ends", async () => {
+    const draftStream = createDraftStream(999);
+    createTelegramDraftStream.mockReturnValue(draftStream);
+    dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
+      async ({ dispatcherOptions, replyOptions }) => {
+        await replyOptions?.onPartialReply?.({ text: "Let me check" });
+        await replyOptions?.onReasoningEnd?.();
+        await replyOptions?.onPartialReply?.({ text: "Here's the answer" });
+        await dispatcherOptions.deliver({ text: "Here's the answer" }, { kind: "final" });
+        return { queuedFinal: true };
+      },
+    );
+    deliverReplies.mockResolvedValue({ delivered: true });
+
+    await dispatchWithContext({ context: createContext(), streamMode: "partial" });
+
+    expect(draftStream.forceNewMessage).not.toHaveBeenCalled();
+  });
+
   it("does not force new message on reasoning end without previous output", async () => {
     const draftStream = createDraftStream(999);
     createTelegramDraftStream.mockReturnValue(draftStream);
diff --git a/src/telegram/bot-message-dispatch.ts b/src/telegram/bot-message-dispatch.ts
index b0ff06b93b2..7cfd0778790 100644
--- a/src/telegram/bot-message-dispatch.ts
+++ b/src/telegram/bot-message-dispatch.ts
@@ -113,6 +113,7 @@ export const dispatchTelegramMessage = async ({
     draftStream && streamMode === "block"
       ? resolveTelegramDraftStreamingChunking(cfg, route.accountId)
       : undefined;
+  const shouldSplitPreviewMessages = streamMode === "block";
   const draftChunker = draftChunking ? new EmbeddedBlockChunker(draftChunking) : undefined;
   const mediaLocalRoots = getAgentScopedMediaLocalRoots(cfg, route.agentId);
   let lastPartialText = "";
@@ -424,13 +425,12 @@ export const dispatchTelegramMessage = async ({
         onPartialReply: draftStream ? (payload) => updateDraftFromPartial(payload.text) : undefined,
         onAssistantMessageStart: draftStream
           ? () => {
-              // When a new assistant message starts (e.g., after tool call),
-              // force a new Telegram message if we have previous content.
-              // Only force once per response to avoid excessive splitting.
+              // Only split preview bubbles in block mode. In partial mode, keep
+              // editing one preview message to avoid flooding the chat.
               logVerbose(
                 `telegram: onAssistantMessageStart called, hasStreamedMessage=${hasStreamedMessage}`,
               );
-              if (hasStreamedMessage) {
+              if (shouldSplitPreviewMessages && hasStreamedMessage) {
                 logVerbose(`telegram: calling forceNewMessage()`);
                 draftStream.forceNewMessage();
               }
@@ -441,13 +441,13 @@ export const dispatchTelegramMessage = async ({
           : undefined,
         onReasoningEnd: draftStream
           ? () => {
-              // When a thinking block ends, force a new Telegram message for the next text output.
-              if (hasStreamedMessage) {
+              // Same policy as assistant-message boundaries: split only in block mode.
+              if (shouldSplitPreviewMessages && hasStreamedMessage) {
                 draftStream.forceNewMessage();
-                lastPartialText = "";
-                draftText = "";
-                draftChunker?.reset();
               }
+              lastPartialText = "";
+              draftText = "";
+              draftChunker?.reset();
             }
           : undefined,
         onModelSelected,