feat: add sendMessageDraft streaming mode (#141)

yaffalhakim1 · web-flow · commit e1d52d66a783 · 2026-06-04T14:23:45.000+03:00
diff --git a/.env.example b/.env.example
@@ -80,6 +80,12 @@ OPENCODE_MODEL_ID=big-pickle
 # If exceeded, the bot stops waiting for the result and marks the run as failed.
 # SCHEDULED_TASK_EXECUTION_TIMEOUT_MINUTES=120
 
+# Response streaming mode: "edit" (default) or "draft" (default: edit)
+# edit = uses editMessageText to update messages incrementally (may flicker)
+# draft = uses sendMessageDraft (Bot API 9.5+) for smooth animated draft previews,
+#         then finalizes with sendMessage. Only works in private chats.
+# RESPONSE_STREAMING_MODE=edit
+
 # Stream update throttle in milliseconds for assistant/tool message edits (default: 1000)
 # Higher value = fewer Telegram edit requests, lower value = more real-time updates
 # RESPONSE_STREAM_THROTTLE_MS=1000
diff --git a/src/bot/index.ts b/src/bot/index.ts
@@ -93,9 +93,11 @@ import { sendTtsResponseForSession } from "./utils/send-tts-response.js";
 import { deliverThinkingMessage } from "./utils/thinking-message.js";
 import { shouldSuppressUserAbortSessionError } from "./utils/abort-error-suppression.js";
 import {
+  completeDraftPart,
   editRenderedBotPart,
   getTelegramRenderedPartSignature,
   sendBotText,
+  sendDraftBotPart,
   sendRenderedBotPart,
 } from "./utils/telegram-text.js";
 import { formatAssistantRunFooter } from "./utils/assistant-run-footer.js";
@@ -134,6 +136,7 @@ let unsubscribeReadyRestore: (() => void) | null = null;
 
 const TELEGRAM_DOCUMENT_CAPTION_MAX_LENGTH = 1024;
 const RESPONSE_STREAM_THROTTLE_MS = config.bot.responseStreamThrottleMs;
+const RESPONSE_STREAMING_MODE = config.bot.responseStreamingMode;
 const RESPONSE_STREAM_TEXT_LIMIT = 3800;
 const SESSION_RETRY_PREFIX = "🔁";
 const SUBAGENT_STREAM_PREFIX = "🧩";
@@ -237,64 +240,116 @@ const toolMessageBatcher = new ToolMessageBatcher({
   },
 });
 
-const responseStreamer = new ResponseStreamer({
-  throttleMs: RESPONSE_STREAM_THROTTLE_MS,
-  sendPart: async (part, options) => {
-    if (!botInstance || !chatIdInstance || chatIdInstance <= 0) {
-      throw new Error("Bot context missing for streamed send");
-    }
+let nextDraftId = 1;
+function getNextDraftId(): number {
+  const id = nextDraftId;
+  nextDraftId += 1;
+  return id;
+}
 
-    return sendRenderedBotPart({
-      api: botInstance.api,
-      chatId: chatIdInstance,
-      part,
-      options,
-    });
-  },
-  editPart: async (messageId, part, options) => {
-    if (!botInstance || !chatIdInstance || chatIdInstance <= 0) {
-      throw new Error("Bot context missing for streamed edit");
-    }
+const responseStreamer = RESPONSE_STREAMING_MODE === "draft"
+  ? new ResponseStreamer({
+      throttleMs: RESPONSE_STREAM_THROTTLE_MS,
+      sendPart: async (part) => {
+        if (!botInstance || !chatIdInstance || chatIdInstance <= 0) {
+          throw new Error("Bot context missing for draft send");
+        }
 
-    try {
-      return await editRenderedBotPart({
-        api: botInstance.api,
-        chatId: chatIdInstance,
-        messageId,
-        part,
-        options,
-      });
-    } catch (error) {
-      const errorMessage =
-        error instanceof Error ? error.message.toLowerCase() : String(error).toLowerCase();
-      if (errorMessage.includes("message is not modified")) {
-        return {
-          deliveredSignature: getTelegramRenderedPartSignature(part),
-        };
-      }
+        const draftId = getNextDraftId();
+        const result = await sendDraftBotPart({
+          api: botInstance.api,
+          chatId: chatIdInstance,
+          draftId,
+          part,
+        });
+        return { messageId: draftId, deliveredSignature: result.deliveredSignature };
+      },
+      editPart: async (messageId, part) => {
+        if (!botInstance || !chatIdInstance || chatIdInstance <= 0) {
+          throw new Error("Bot context missing for draft edit");
+        }
 
-      throw error;
-    }
-  },
-  deleteText: async (messageId) => {
-    if (!botInstance || !chatIdInstance || chatIdInstance <= 0) {
-      throw new Error("Bot context missing for streamed delete");
-    }
+        return sendDraftBotPart({
+          api: botInstance.api,
+          chatId: chatIdInstance,
+          draftId: messageId,
+          part,
+        });
+      },
+      deleteText: async () => {
+        // Drafts are ephemeral — no need to delete
+      },
+      completePart: async (part, options) => {
+        if (!botInstance || !chatIdInstance || chatIdInstance <= 0) {
+          throw new Error("Bot context missing for draft complete");
+        }
 
-    await botInstance.api.deleteMessage(chatIdInstance, messageId).catch((error) => {
-      const errorMessage =
-        error instanceof Error ? error.message.toLowerCase() : String(error).toLowerCase();
-      if (
-        errorMessage.includes("message to delete not found") ||
-        errorMessage.includes("message identifier is not specified")
-      ) {
-        return;
-      }
+        return completeDraftPart({
+          api: botInstance.api,
+          chatId: chatIdInstance,
+          part,
+          options,
+        });
+      },
+    })
+  : new ResponseStreamer({
+      throttleMs: RESPONSE_STREAM_THROTTLE_MS,
+      sendPart: async (part, options) => {
+        if (!botInstance || !chatIdInstance || chatIdInstance <= 0) {
+          throw new Error("Bot context missing for streamed send");
+        }
 
-      throw error;
+        return sendRenderedBotPart({
+          api: botInstance.api,
+          chatId: chatIdInstance,
+          part,
+          options,
+        });
+      },
+      editPart: async (messageId, part, options) => {
+        if (!botInstance || !chatIdInstance || chatIdInstance <= 0) {
+          throw new Error("Bot context missing for streamed edit");
+        }
+
+        try {
+          return await editRenderedBotPart({
+            api: botInstance.api,
+            chatId: chatIdInstance,
+            messageId,
+            part,
+            options,
+          });
+        } catch (error) {
+          const errorMessage =
+            error instanceof Error ? error.message.toLowerCase() : String(error).toLowerCase();
+          if (errorMessage.includes("message is not modified")) {
+            return {
+              deliveredSignature: getTelegramRenderedPartSignature(part),
+            };
+          }
+
+          throw error;
+        }
+      },
+      deleteText: async (messageId) => {
+        if (!botInstance || !chatIdInstance || chatIdInstance <= 0) {
+          throw new Error("Bot context missing for streamed delete");
+        }
+
+        await botInstance.api.deleteMessage(chatIdInstance, messageId).catch((error) => {
+          const errorMessage =
+            error instanceof Error ? error.message.toLowerCase() : String(error).toLowerCase();
+          if (
+            errorMessage.includes("message to delete not found") ||
+            errorMessage.includes("message identifier is not specified")
+          ) {
+            return;
+          }
+
+          throw error;
+        });
+      },
     });
-  },
-});
 
 setResponseStreamerForReconciliation(responseStreamer);
 
diff --git a/src/bot/streaming/response-streamer.ts b/src/bot/streaming/response-streamer.ts
@@ -35,6 +35,10 @@ interface ResponseStreamerOptions {
     options?: TelegramEditMessageOptions,
   ) => Promise<{ deliveredSignature: string }>;
   deleteText: (messageId: number) => Promise<void>;
+  completePart?: (
+    part: TelegramRenderedPart,
+    options?: TelegramSendMessageOptions,
+  ) => Promise<{ messageId: number; deliveredSignature: string }>;
 }
 
 interface StreamState {
@@ -121,13 +125,15 @@ export class ResponseStreamer {
   private readonly sendPart: ResponseStreamerOptions["sendPart"];
   private readonly editPart: ResponseStreamerOptions["editPart"];
   private readonly deleteText: ResponseStreamerOptions["deleteText"];
+  private readonly completePart: ResponseStreamerOptions["completePart"];
   private readonly states: Map<string, StreamState> = new Map();
 
   constructor(options: ResponseStreamerOptions) {
     this.throttleMs = Math.max(0, Math.floor(options.throttleMs));
     this.sendPart = options.sendPart;
     this.editPart = options.editPart;
     this.deleteText = options.deleteText;
+    this.completePart = options.completePart;
   }
 
   enqueue(sessionId: string, messageId: string, payload: StreamingMessagePayload): void {
@@ -191,8 +197,25 @@ export class ResponseStreamer {
     let synced = true;
     if (options?.flushFinal !== false) {
       synced = await this.enqueueTask(state, () => this.flushState(state, "complete"));
-      if (!synced && state.isBroken) {
-        await this.cleanupBrokenStream(state, "final_sync_failed_cleanup");
+    }
+
+    if (synced && this.completePart && state.latestPayload) {
+      try {
+        const realMessageIds: number[] = [];
+        for (const part of state.latestPayload.parts) {
+          if (!part.text) {
+            continue;
+          }
+          const result = await this.completePart(part, state.latestPayload.sendOptions);
+          realMessageIds.push(result.messageId);
+        }
+        state.telegramMessageIds = realMessageIds;
+      } catch (error) {
+        logger.error(
+          `[ResponseStreamer] Failed to persist draft message: session=${sessionId}, message=${messageId}`,
+          error,
+        );
+        synced = false;
       }
     }
 
diff --git a/src/bot/utils/telegram-text.ts b/src/bot/utils/telegram-text.ts
@@ -8,6 +8,7 @@ import type { TelegramRenderedPart } from "../../telegram/render/types.js";
 
 type SendMessageApi = Pick<Api<RawApi>, "sendMessage">;
 type EditMessageApi = Pick<Api<RawApi>, "editMessageText">;
+type SendDraftApi = Pick<Api<RawApi>, "sendMessageDraft">;
 
 type TelegramSendMessageOptions = Parameters<SendMessageApi["sendMessage"]>[2];
 type TelegramEditMessageOptions = Parameters<EditMessageApi["editMessageText"]>[3];
@@ -203,6 +204,95 @@ export async function editRenderedBotPart({
   }
 }
 
+interface SendDraftBotPartParams {
+  api: SendDraftApi;
+  chatId: Parameters<SendDraftApi["sendMessageDraft"]>[0];
+  draftId: number;
+  part: TelegramRenderedPart;
+}
+
+interface CompleteDraftPartParams {
+  api: SendMessageApi;
+  chatId: Parameters<SendMessageApi["sendMessage"]>[0];
+  part: TelegramRenderedPart;
+  options?: TelegramSendMessageOptions;
+}
+
+export async function sendDraftBotPart({
+  api,
+  chatId,
+  draftId,
+  part,
+}: SendDraftBotPartParams): Promise<RenderedPartDeliveryResult> {
+  logger.debug("[Bot] Sending draft part", {
+    draftId,
+    textLength: part.text.length,
+    entityCount: part.entities?.length ?? 0,
+  });
+
+  if (!part.entities?.length) {
+    await api.sendMessageDraft(chatId, draftId, part.text);
+    return {
+      deliveredSignature: getTelegramRenderedPartSignature({ text: part.text }),
+    };
+  }
+
+  try {
+    await api.sendMessageDraft(chatId, draftId, part.text, {
+      entities: part.entities,
+    });
+    return {
+      deliveredSignature: getTelegramRenderedPartSignature(part),
+    };
+  } catch (error) {
+    logger.warn("[Bot] Entity draft failed, retrying with raw fallback", error);
+    await api.sendMessageDraft(chatId, draftId, part.fallbackText);
+    return {
+      deliveredSignature: getTelegramRenderedPartSignature({ text: part.fallbackText }),
+    };
+  }
+}
+
+export async function completeDraftPart({
+  api,
+  chatId,
+  part,
+  options,
+}: CompleteDraftPartParams): Promise<RenderedPartSendResult> {
+  const rawOptions = stripRichFormattingOptions(options);
+
+  logger.debug("[Bot] Completing draft with real message", {
+    textLength: part.text.length,
+    entityCount: part.entities?.length ?? 0,
+  });
+
+  if (!part.entities?.length) {
+    const sentMessage = await api.sendMessage(chatId, part.text, rawOptions);
+    return {
+      messageId: sentMessage.message_id,
+      deliveredSignature: getTelegramRenderedPartSignature({ text: part.text }),
+    };
+  }
+
+  try {
+    const sentMessage = await api.sendMessage(chatId, part.text, {
+      ...(rawOptions || {}),
+      entities: part.entities,
+    });
+    return {
+      messageId: sentMessage.message_id,
+      deliveredSignature: getTelegramRenderedPartSignature(part),
+    };
+  } catch (error) {
+    logger.warn("[Bot] Entity complete failed, retrying with raw fallback", error);
+    const sentMessage = await api.sendMessage(chatId, part.fallbackText, rawOptions);
+    return {
+      messageId: sentMessage.message_id,
+      deliveredSignature: getTelegramRenderedPartSignature({ text: part.fallbackText }),
+    };
+  }
+}
+
 export async function editBotText({
   api,
   chatId,
diff --git a/src/config.ts b/src/config.ts
diff --git a/tests/bot/streaming/response-streamer.test.ts b/tests/bot/streaming/response-streamer.test.ts
