fix: centralize draft preview finalization

extensions/matrix/src/matrix/draft-stream.test.ts

@@ -182,6 +182,7 @@ describe("createMatrixDraftStream", () => {
       .mockReset()
       .mockImplementation((text: string) => (text ? [text] : []));
     convertMarkdownTablesMock.mockReset().mockImplementation((text: string) => text);
+    sendModuleMocks.editMessageMatrix.mockClear();
   });
 
   afterEach(() => {
@@ -503,6 +504,24 @@ describe("createMatrixDraftStream", () => {
     );
   });
 
+  it("discardPending cancels pending updates without creating another preview event", async () => {
+    const stream = createMatrixDraftStream({
+      roomId: "!room:test",
+      client,
+      cfg: {} as import("../types.js").CoreConfig,
+    });
+
+    stream.update("First draft");
+    await stream.flush();
+    stream.update("Pending draft");
+    await stream.discardPending();
+    await stream.flush();
+
+    expect(sendMessageMock).toHaveBeenCalledTimes(1);
+    expect(sendModuleMocks.editMessageMatrix).not.toHaveBeenCalled();
+    expect(stream.eventId()).toBe("$evt1");
+  });
+
   it("uses converted Matrix text when checking the single-event preview limit", async () => {
     const log = vi.fn();
     resolveTextChunkLimitMock.mockReturnValue(5);

extensions/matrix/src/matrix/draft-stream.ts

@@ -29,6 +29,8 @@ export type MatrixDraftStream = {
   flush: () => Promise<void>;
   /** Flush and mark this block as done. Returns the event ID if a message was sent. */
   stop: () => Promise<string | undefined>;
+  /** Cancel pending draft updates without creating a new preview event. */
+  discardPending: () => Promise<void>;
   /** Clear the MSC4357 live marker in place when the draft is kept as final text. */
   finalizeLive: () => Promise<boolean>;
   /** Reset state for the next text block (after tool calls). */
@@ -180,6 +182,12 @@ export function createMatrixDraftStream(params: {
     return currentEventId;
   };
 
+  const discardPending = async (): Promise<void> => {
+    stopped = true;
+    loop.stop();
+    await loop.waitForInFlight();
+  };
+
   const reset = (): void => {
     // Clear reply context unless preserveReplyId is set (replyToMode "all"),
     // in which case subsequent blocks should keep replying to the original.
@@ -203,6 +211,7 @@ export function createMatrixDraftStream(params: {
     },
     flush: loop.flush,
     stop,
+    discardPending,
     finalizeLive,
     reset,
     eventId: () => currentEventId,

extensions/matrix/src/matrix/monitor/handler.test.ts

@@ -3231,6 +3231,50 @@ describe("matrix monitor handler draft streaming", () => {
     await finish();
   });
 
+  it("does not create a throwaway draft for fast media-only finals", async () => {
+    const { dispatch, redactEventMock } = createStreamingHarness();
+    const { deliver, finish } = await dispatch();
+
+    await deliver({ mediaUrl: "https://example.com/image.png" }, { kind: "final" });
+
+    expect(sendSingleTextMessageMatrixMock).not.toHaveBeenCalled();
+    expect(editMessageMatrixMock).not.toHaveBeenCalled();
+    expect(redactEventMock).not.toHaveBeenCalled();
+    expect(deliverMatrixRepliesMock).toHaveBeenCalledTimes(1);
+    await finish();
+  });
+
+  it("does not create a throwaway draft for fast error finals", async () => {
+    const { dispatch, redactEventMock } = createStreamingHarness();
+    const { deliver, finish } = await dispatch();
+
+    await deliver({ text: "Something failed", isError: true } as never, { kind: "final" });
+
+    expect(sendSingleTextMessageMatrixMock).not.toHaveBeenCalled();
+    expect(editMessageMatrixMock).not.toHaveBeenCalled();
+    expect(redactEventMock).not.toHaveBeenCalled();
+    expect(deliverMatrixRepliesMock).toHaveBeenCalledTimes(1);
+    await finish();
+  });
+
+  it("redacts existing drafts for text error finals and uses normal delivery", async () => {
+    const { dispatch, redactEventMock } = createStreamingHarness();
+    const { deliver, opts, finish } = await dispatch();
+
+    opts.onPartialReply?.({ text: "Partial reply" });
+    await vi.waitFor(() => {
+      expect(sendSingleTextMessageMatrixMock).toHaveBeenCalledTimes(1);
+    });
+
+    deliverMatrixRepliesMock.mockClear();
+    await deliver({ text: "Something failed", isError: true } as never, { kind: "final" });
+
+    expect(editMessageMatrixMock).not.toHaveBeenCalled();
+    expect(redactEventMock).toHaveBeenCalledWith("!room:example.org", "$draft1");
+    expect(deliverMatrixRepliesMock).toHaveBeenCalledTimes(1);
+    await finish();
+  });
+
   it("finalizes partial drafts before reusing unchanged media captions", async () => {
     const { dispatch, redactEventMock } = createStreamingHarness({ streaming: "partial" });
     const { deliver, opts, finish } = await dispatch();

extensions/matrix/src/matrix/monitor/handler.ts

@@ -121,6 +121,7 @@ type MatrixAllowBotsMode = "off" | "mentions" | "all";
 type MatrixDraftStreamHandle = {
   update: (text: string) => void;
   stop: () => Promise<string | undefined>;
+  discardPending: () => Promise<void>;
   eventId: () => string | undefined;
   mustDeliverFinalNormally: () => boolean;
   matchesPreparedText: (text: string) => boolean;
@@ -1547,10 +1548,8 @@ export function createMatrixRoomMessageHandler(params: MatrixMonitorHandlerParam
             if (draftStream && info.kind !== "tool" && !payload.isCompactionNotice) {
               const hasMedia = Boolean(payload.mediaUrl) || (payload.mediaUrls?.length ?? 0) > 0;
 
-              await draftStream.stop();
-              const draftEventId = draftStream.eventId();
-
               if (draftConsumed) {
+                await draftStream.discardPending();
                 await deliverMatrixReplies({
                   cfg,
                   replies: [payload],
@@ -1572,11 +1571,25 @@ export function createMatrixRoomMessageHandler(params: MatrixMonitorHandlerParam
                 replyToMode !== "off" &&
                 !threadTarget &&
                 payloadReplyToId !== currentDraftReplyToId;
-              const mustDeliverFinalNormally = draftStream.mustDeliverFinalNormally();
+              let mustDeliverFinalNormally = draftStream.mustDeliverFinalNormally();
+              const canPotentiallyFinalizeDraft =
+                Boolean(payload.text?.trim()) &&
+                !payload.isError &&
+                !payloadReplyMismatch &&
+                !mustDeliverFinalNormally;
+
+              if (canPotentiallyFinalizeDraft) {
+                await draftStream.stop();
+                mustDeliverFinalNormally = draftStream.mustDeliverFinalNormally();
+              } else {
+                await draftStream.discardPending();
+              }
+              const draftEventId = draftStream.eventId();
 
               if (
                 draftEventId &&
                 payload.text &&
+                !payload.isError &&
                 !hasMedia &&
                 !payloadReplyMismatch &&
                 !mustDeliverFinalNormally
@@ -1666,7 +1679,8 @@ export function createMatrixRoomMessageHandler(params: MatrixMonitorHandlerParam
                 draftConsumed = true;
               } else {
                 const draftRedacted =
-                  Boolean(draftEventId) && (payloadReplyMismatch || mustDeliverFinalNormally);
+                  Boolean(draftEventId) &&
+                  (payload.isError || payloadReplyMismatch || mustDeliverFinalNormally);
                 if (draftRedacted && draftEventId) {
                   await redactMatrixDraftEvent(client, roomId, draftEventId);
                 }