From 1f822d7c222cb9671d958d38f775f680d8045f82 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Thu, 7 May 2026 01:58:00 +0100
Subject: [PATCH] test(telegram): lock draft finalization ordering

---
 .../telegram/src/bot-message-dispatch.test.ts | 18 ++++++
 src/channels/draft-stream-controls.test.ts    | 62 +++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/extensions/telegram/src/bot-message-dispatch.test.ts b/extensions/telegram/src/bot-message-dispatch.test.ts
index b634ec9cb8f..9b6c9d2b35a 100644
--- a/extensions/telegram/src/bot-message-dispatch.test.ts
+++ b/extensions/telegram/src/bot-message-dispatch.test.ts
@@ -924,6 +924,24 @@ describe("dispatchTelegramMessage draft streaming", () => {
     expect(deliverReplies).not.toHaveBeenCalled();
   });
 
+  it("waits for queued draft-lane partials before finalizing the Telegram reply", async () => {
+    const { answerDraftStream } = setupDraftStreams({ answerMessageId: 2001 });
+    dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
+      async ({ dispatcherOptions, replyOptions }) => {
+        const pendingPartial = replyOptions?.onPartialReply?.({ text: "Working" });
+        await dispatcherOptions.deliver({ text: "Done" }, { kind: "final" });
+        await pendingPartial;
+        return { queuedFinal: true };
+      },
+    );
+
+    await dispatchWithContext({ context: createContext() });
+
+    expect(answerDraftStream.update).toHaveBeenNthCalledWith(1, "Working");
+    expect(answerDraftStream.update).toHaveBeenNthCalledWith(2, "Done");
+    expect(deliverReplies).not.toHaveBeenCalled();
+  });
+
   it("keeps progress updates in a draft and sends the final answer normally", async () => {
     const { answerDraftStream } = setupDraftStreams({ answerMessageId: 2001 });
     dispatchReplyWithBufferedBlockDispatcher.mockImplementation(
diff --git a/src/channels/draft-stream-controls.test.ts b/src/channels/draft-stream-controls.test.ts
index 4c627aff636..9089fc485ce 100644
--- a/src/channels/draft-stream-controls.test.ts
+++ b/src/channels/draft-stream-controls.test.ts
@@ -120,6 +120,68 @@ describe("draft-stream-controls", () => {
     expect(deleteMessage).toHaveBeenCalledWith("m-4");
   });
 
+  it("lifecycle clear cancels pending draft text instead of flushing it", async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(0);
+    try {
+      const state = { stopped: false, final: false };
+      let messageId: string | undefined = "m-6";
+      const sendOrEditStreamMessage = vi.fn(async () => true);
+      const deleteMessage = vi.fn(async () => {});
+
+      const lifecycle = createFinalizableDraftLifecycle({
+        throttleMs: 250,
+        state,
+        sendOrEditStreamMessage,
+        readMessageId: () => messageId,
+        clearMessageId: () => {
+          messageId = undefined;
+        },
+        isValidMessageId: (value): value is string => typeof value === "string",
+        deleteMessage,
+        warnPrefix: "cleanup failed",
+      });
+
+      lifecycle.update("pending draft");
+      await lifecycle.clear();
+
+      expect(state.stopped).toBe(true);
+      expect(sendOrEditStreamMessage).not.toHaveBeenCalled();
+      expect(deleteMessage).toHaveBeenCalledWith("m-6");
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("lifecycle stop flushes pending final draft text", async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(0);
+    try {
+      const state = { stopped: false, final: false };
+      const sendOrEditStreamMessage = vi.fn(async () => true);
+
+      const lifecycle = createFinalizableDraftLifecycle({
+        throttleMs: 250,
+        state,
+        sendOrEditStreamMessage,
+        readMessageId: () => "m-7",
+        clearMessageId: () => {},
+        isValidMessageId: (value): value is string => typeof value === "string",
+        deleteMessage: async () => {},
+        warnPrefix: "cleanup failed",
+      });
+
+      lifecycle.update("final draft");
+      await lifecycle.stop();
+
+      expect(state.final).toBe(true);
+      expect(sendOrEditStreamMessage).toHaveBeenCalledTimes(1);
+      expect(sendOrEditStreamMessage).toHaveBeenCalledWith("final draft");
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
   it("lifecycle seal ignores late updates without clearing the preview id", async () => {
     const state = { stopped: false, final: false };
     let messageId: string | undefined = "m-5";