From e45fcc57ed3aa5febf6ea525b4795505ce4c1668 Mon Sep 17 00:00:00 2001
From: Ayaan Zaidi <zaidi@uplause.io>
Date: Sun, 8 Mar 2026 08:18:01 +0530
Subject: [PATCH] fix(telegram): restore DM draft streaming

---
 src/telegram/bot-message-dispatch.test.ts | 35 ++++++++++++++++++++---
 src/telegram/bot-message-dispatch.ts      |  9 ++----
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/src/telegram/bot-message-dispatch.test.ts b/src/telegram/bot-message-dispatch.test.ts
index ddec14d60e2..1e8202bce67 100644
--- a/src/telegram/bot-message-dispatch.test.ts
+++ b/src/telegram/bot-message-dispatch.test.ts
@@ -1171,7 +1171,7 @@ describe("dispatchTelegramMessage draft streaming", () => {
     },
   );
 
-  it("uses message preview transport for all DM lanes when streaming is active", async () => {
+  it("uses message preview transport for DM reasoning lane when answer preview lane is active", async () => {
     setupDraftStreams({ answerMessageId: 999, reasoningMessageId: 111 });
     dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
       async ({ dispatcherOptions, replyOptions }) => {
@@ -1187,12 +1187,10 @@ describe("dispatchTelegramMessage draft streaming", () => {
     await dispatchWithContext({ context: createReasoningStreamContext(), streamMode: "partial" });
 
     expect(createTelegramDraftStream).toHaveBeenCalledTimes(2);
-    // Both answer (call[0]) and reasoning (call[1]) lanes should use message
-    // transport in DMs to prevent duplicate messages. (Fixes #33453)
     expect(createTelegramDraftStream.mock.calls[0]?.[0]).toEqual(
       expect.objectContaining({
         thread: { id: 777, scope: "dm" },
-        previewTransport: "message",
+        previewTransport: "auto",
       }),
     );
     expect(createTelegramDraftStream.mock.calls[1]?.[0]).toEqual(
@@ -1203,6 +1201,35 @@ describe("dispatchTelegramMessage draft streaming", () => {
     );
   });
 
+  it("materializes DM answer draft final without sending a duplicate final message", async () => {
+    const answerDraftStream = createTestDraftStream({ previewMode: "draft" });
+    answerDraftStream.materialize.mockResolvedValue(321);
+    const reasoningDraftStream = createDraftStream(111);
+    createTelegramDraftStream
+      .mockImplementationOnce(() => answerDraftStream)
+      .mockImplementationOnce(() => reasoningDraftStream);
+    dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
+      async ({ dispatcherOptions, replyOptions }) => {
+        await replyOptions?.onPartialReply?.({ text: "Checking the directory..." });
+        await dispatcherOptions.deliver({ text: "Checking the directory..." }, { kind: "final" });
+        return { queuedFinal: true };
+      },
+    );
+    deliverReplies.mockResolvedValue({ delivered: true });
+
+    await dispatchWithContext({ context: createContext(), streamMode: "partial" });
+
+    expect(createTelegramDraftStream.mock.calls[0]?.[0]).toEqual(
+      expect.objectContaining({
+        thread: { id: 777, scope: "dm" },
+        previewTransport: "auto",
+      }),
+    );
+    expect(answerDraftStream.materialize).toHaveBeenCalledTimes(1);
+    expect(deliverReplies).not.toHaveBeenCalled();
+    expect(editMessageTelegram).not.toHaveBeenCalled();
+  });
+
   it("keeps reasoning and answer streaming in separate preview lanes", async () => {
     const { answerDraftStream, reasoningDraftStream } = setupDraftStreams({
       answerMessageId: 999,
diff --git a/src/telegram/bot-message-dispatch.ts b/src/telegram/bot-message-dispatch.ts
index 111e5991a72..859a35688f6 100644
--- a/src/telegram/bot-message-dispatch.ts
+++ b/src/telegram/bot-message-dispatch.ts
@@ -190,22 +190,19 @@ export const dispatchTelegramMessage = async ({
   const draftReplyToMessageId =
     replyToMode !== "off" && typeof msg.message_id === "number" ? msg.message_id : undefined;
   const draftMinInitialChars = DRAFT_MIN_INITIAL_CHARS;
-  // Use message transport (sendMessage + editMessageText) for all lanes in
-  // DMs so that streamMessageId is tracked. Draft transport doesn't track a
-  // messageId, causing resolvePreviewTarget() to miss the preview on final
-  // delivery — which sends a duplicate message. (Fixes #33453)
-  const useMessagePreviewTransportForDm = threadSpec?.scope === "dm" && canStreamAnswerDraft;
   const mediaLocalRoots = getAgentScopedMediaLocalRoots(cfg, route.agentId);
   const archivedAnswerPreviews: ArchivedPreview[] = [];
   const archivedReasoningPreviewIds: number[] = [];
   const createDraftLane = (laneName: LaneName, enabled: boolean): DraftLaneState => {
+    const useMessagePreviewTransportForDmReasoning =
+      laneName === "reasoning" && threadSpec?.scope === "dm" && canStreamAnswerDraft;
     const stream = enabled
       ? createTelegramDraftStream({
           api: bot.api,
           chatId,
           maxChars: draftMaxChars,
           thread: threadSpec,
-          previewTransport: useMessagePreviewTransportForDm ? "message" : "auto",
+          previewTransport: useMessagePreviewTransportForDmReasoning ? "message" : "auto",
           replyToMessageId: draftReplyToMessageId,
           minInitialChars: draftMinInitialChars,
           renderText: renderDraftPreview,