From 68bb76519ad4677dfa2257ea5938486d4fe26870 Mon Sep 17 00:00:00 2001
From: Gustavo Madeira Santana <gumadeiras@gmail.com>
Date: Thu, 2 Apr 2026 03:47:57 -0400
Subject: [PATCH] Matrix: fix delayed draft block boundaries

---
 .../matrix/src/matrix/monitor/handler.test.ts | 44 +++++++++++++++++++
 .../matrix/src/matrix/monitor/handler.ts      | 11 ++---
 2 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/extensions/matrix/src/matrix/monitor/handler.test.ts b/extensions/matrix/src/matrix/monitor/handler.test.ts
index abc0a8168b1..15b7e829f98 100644
--- a/extensions/matrix/src/matrix/monitor/handler.test.ts
+++ b/extensions/matrix/src/matrix/monitor/handler.test.ts
@@ -1729,6 +1729,50 @@ describe("matrix monitor handler draft streaming", () => {
     await finish();
   });
 
+  it("keeps delayed same-message block boundaries at the emitted block length", async () => {
+    const { dispatch, redactEventMock } = createStreamingHarness({ blockStreamingEnabled: true });
+    const { deliver, opts, finish } = await dispatch();
+
+    opts.onPartialReply?.({ text: "Alpha" });
+    await vi.waitFor(() => {
+      expect(sendSingleTextMessageMatrixMock).toHaveBeenCalledTimes(1);
+    });
+
+    opts.onPartialReply?.({ text: "AlphaBeta" });
+    await vi.waitFor(() => {
+      expect(editMessageMatrixMock).toHaveBeenCalledWith(
+        "!room:example.org",
+        "$draft1",
+        "AlphaBeta",
+        expect.anything(),
+      );
+    });
+
+    await opts.onBlockReplyQueued?.({ text: "Alpha" });
+
+    sendSingleTextMessageMatrixMock.mockClear();
+    editMessageMatrixMock.mockClear();
+    sendSingleTextMessageMatrixMock.mockResolvedValueOnce({
+      messageId: "$draft2",
+      roomId: "!room",
+    });
+    await deliver({ text: "Alpha" }, { kind: "block" });
+
+    await vi.waitFor(() => {
+      expect(sendSingleTextMessageMatrixMock).toHaveBeenCalledTimes(1);
+    });
+    expect(sendSingleTextMessageMatrixMock.mock.calls[0]?.[1]).toBe("Beta");
+    expect(editMessageMatrixMock).toHaveBeenCalledWith(
+      "!room:example.org",
+      "$draft1",
+      "Alpha",
+      expect.anything(),
+    );
+    expect(deliverMatrixRepliesMock).not.toHaveBeenCalled();
+    expect(redactEventMock).not.toHaveBeenCalled();
+    await finish();
+  });
+
   it("falls back to deliverMatrixReplies when final edit fails", async () => {
     const { dispatch } = createStreamingHarness();
     const { deliver, opts, finish } = await dispatch();
diff --git a/extensions/matrix/src/matrix/monitor/handler.ts b/extensions/matrix/src/matrix/monitor/handler.ts
index 7899240c64a..c4002827a2a 100644
--- a/extensions/matrix/src/matrix/monitor/handler.ts
+++ b/extensions/matrix/src/matrix/monitor/handler.ts
@@ -1184,13 +1184,10 @@ export function createMatrixRoomMessageHandler(params: MatrixMonitorHandlerParam
         const messageGeneration = context?.assistantMessageIndex ?? currentDraftMessageGeneration;
         const lastQueuedDraftBoundaryOffset =
           latestQueuedDraftBoundaryOffsets.get(messageGeneration) ?? 0;
-        const nextDraftBoundaryOffset =
-          messageGeneration === currentDraftMessageGeneration
-            ? Math.max(
-                latestDraftFullText.length,
-                lastQueuedDraftBoundaryOffset + payloadTextLength,
-              )
-            : lastQueuedDraftBoundaryOffset + payloadTextLength;
+        // Logical block boundaries must follow emitted block text, not whichever
+        // later partial preview has already arrived by the time the async
+        // boundary callback drains.
+        const nextDraftBoundaryOffset = lastQueuedDraftBoundaryOffset + payloadTextLength;
         latestQueuedDraftBoundaryOffsets.set(messageGeneration, nextDraftBoundaryOffset);
         pendingDraftBoundaries.push({
           messageGeneration,