feat(telegram): use sendMessageDraft for private chat streaming (#31824)

* feat(telegram): use sendMessageDraft for private stream previews

* test(telegram): cover DM draft id rotation race

* fix(telegram): keep DM reasoning updates in draft preview

* fix(telegram): split DM reasoning preview transport

* fix(telegram): harden DM draft preview fallback paths

* style(telegram): normalize draft preview formatting

src/telegram/draft-stream.ts

@@ -4,11 +4,41 @@ import { buildTelegramThreadParams, type TelegramThreadSpec } from "./bot/helper
 
 const TELEGRAM_STREAM_MAX_CHARS = 4096;
 const DEFAULT_THROTTLE_MS = 1000;
+const TELEGRAM_DRAFT_ID_MAX = 2_147_483_647;
+const THREAD_NOT_FOUND_RE = /400:\s*Bad Request:\s*message thread not found/i;
+
+type TelegramSendMessageDraft = (
+  chatId: number,
+  draftId: number,
+  text: string,
+  params?: {
+    message_thread_id?: number;
+    parse_mode?: "HTML";
+  },
+) => Promise<unknown>;
+
+let nextDraftId = 0;
+
+function allocateTelegramDraftId(): number {
+  nextDraftId = nextDraftId >= TELEGRAM_DRAFT_ID_MAX ? 1 : nextDraftId + 1;
+  return nextDraftId;
+}
+
+function resolveSendMessageDraftApi(api: Bot["api"]): TelegramSendMessageDraft | undefined {
+  const sendMessageDraft = (api as Bot["api"] & { sendMessageDraft?: TelegramSendMessageDraft })
+    .sendMessageDraft;
+  if (typeof sendMessageDraft !== "function") {
+    return undefined;
+  }
+  return sendMessageDraft.bind(api as object);
+}
 
 export type TelegramDraftStream = {
   update: (text: string) => void;
   flush: () => Promise<void>;
   messageId: () => number | undefined;
+  previewMode?: () => "message" | "draft";
+  previewRevision?: () => number;
   clear: () => Promise<void>;
   stop: () => Promise<void>;
   /** Reset internal state so the next update creates a new message instead of editing. */
@@ -31,6 +61,7 @@ export function createTelegramDraftStream(params: {
   chatId: number;
   maxChars?: number;
   thread?: TelegramThreadSpec | null;
+  previewTransport?: "auto" | "message" | "draft";
   replyToMessageId?: number;
   throttleMs?: number;
   /** Minimum chars before sending first message (debounce for push notifications) */
@@ -49,17 +80,126 @@ export function createTelegramDraftStream(params: {
   const throttleMs = Math.max(250, params.throttleMs ?? DEFAULT_THROTTLE_MS);
   const minInitialChars = params.minInitialChars;
   const chatId = params.chatId;
+  const requestedPreviewTransport = params.previewTransport ?? "auto";
+  const prefersDraftTransport =
+    requestedPreviewTransport === "draft"
+      ? true
+      : requestedPreviewTransport === "message"
+        ? false
+        : params.thread?.scope === "dm";
   const threadParams = buildTelegramThreadParams(params.thread);
   const replyParams =
     params.replyToMessageId != null
       ? { ...threadParams, reply_to_message_id: params.replyToMessageId }
       : threadParams;
+  const resolvedDraftApi = prefersDraftTransport
+    ? resolveSendMessageDraftApi(params.api)
+    : undefined;
+  const usesDraftTransport = Boolean(prefersDraftTransport && resolvedDraftApi);
+  if (prefersDraftTransport && !usesDraftTransport) {
+    params.warn?.(
+      "telegram stream preview: sendMessageDraft unavailable; falling back to sendMessage/editMessageText",
+    );
+  }
 
   const streamState = { stopped: false, final: false };
   let streamMessageId: number | undefined;
+  let streamDraftId = usesDraftTransport ? allocateTelegramDraftId() : undefined;
   let lastSentText = "";
   let lastSentParseMode: "HTML" | undefined;
+  let previewRevision = 0;
   let generation = 0;
+  const sendStreamPreview = usesDraftTransport
+    ? async ({
+        renderedText,
+        renderedParseMode,
+      }: {
+        renderedText: string;
+        renderedParseMode: "HTML" | undefined;
+        sendGeneration: number;
+      }): Promise<boolean> => {
+        const draftId = streamDraftId ?? allocateTelegramDraftId();
+        streamDraftId = draftId;
+        const draftParams = {
+          ...(threadParams?.message_thread_id != null
+            ? { message_thread_id: threadParams.message_thread_id }
+            : {}),
+          ...(renderedParseMode ? { parse_mode: renderedParseMode } : {}),
+        };
+        await resolvedDraftApi!(
+          chatId,
+          draftId,
+          renderedText,
+          Object.keys(draftParams).length > 0 ? draftParams : undefined,
+        );
+        return true;
+      }
+    : async ({
+        renderedText,
+        renderedParseMode,
+        sendGeneration,
+      }: {
+        renderedText: string;
+        renderedParseMode: "HTML" | undefined;
+        sendGeneration: number;
+      }): Promise<boolean> => {
+        if (typeof streamMessageId === "number") {
+          if (renderedParseMode) {
+            await params.api.editMessageText(chatId, streamMessageId, renderedText, {
+              parse_mode: renderedParseMode,
+            });
+          } else {
+            await params.api.editMessageText(chatId, streamMessageId, renderedText);
+          }
+          return true;
+        }
+        const sendParams = renderedParseMode
+          ? {
+              ...replyParams,
+              parse_mode: renderedParseMode,
+            }
+          : replyParams;
+        let sent;
+        try {
+          sent = await params.api.sendMessage(chatId, renderedText, sendParams);
+        } catch (err) {
+          const hasThreadParam =
+            "message_thread_id" in (sendParams ?? {}) &&
+            typeof (sendParams as { message_thread_id?: unknown }).message_thread_id === "number";
+          if (!hasThreadParam || !THREAD_NOT_FOUND_RE.test(String(err))) {
+            throw err;
+          }
+          const threadlessParams = {
+            ...(sendParams as Record<string, unknown>),
+          };
+          delete threadlessParams.message_thread_id;
+          params.warn?.(
+            "telegram stream preview send failed with message_thread_id, retrying without thread",
+          );
+          sent = await params.api.sendMessage(
+            chatId,
+            renderedText,
+            Object.keys(threadlessParams).length > 0 ? threadlessParams : undefined,
+          );
+        }
+        const sentMessageId = sent?.message_id;
+        if (typeof sentMessageId !== "number" || !Number.isFinite(sentMessageId)) {
+          streamState.stopped = true;
+          params.warn?.("telegram stream preview stopped (missing message id from sendMessage)");
+          return false;
+        }
+        const normalizedMessageId = Math.trunc(sentMessageId);
+        if (sendGeneration !== generation) {
+          params.onSupersededPreview?.({
+            messageId: normalizedMessageId,
+            textSnapshot: renderedText,
+            parseMode: renderedParseMode,
+          });
+          return true;
+        }
+        streamMessageId = normalizedMessageId;
+        return true;
+      };
 
   const sendOrEditStreamMessage = async (text: string): Promise<boolean> => {
     // Allow final flush even if stopped (e.g., after clear()).
@@ -100,40 +240,15 @@ export function createTelegramDraftStream(params: {
     lastSentText = renderedText;
     lastSentParseMode = renderedParseMode;
     try {
-      if (typeof streamMessageId === "number") {
-        if (renderedParseMode) {
-          await params.api.editMessageText(chatId, streamMessageId, renderedText, {
-            parse_mode: renderedParseMode,
-          });
-        } else {
-          await params.api.editMessageText(chatId, streamMessageId, renderedText);
-        }
-        return true;
+      const sent = await sendStreamPreview({
+        renderedText,
+        renderedParseMode,
+        sendGeneration,
+      });
+      if (sent) {
+        previewRevision += 1;
       }
-      const sendParams = renderedParseMode
-        ? {
-            ...replyParams,
-            parse_mode: renderedParseMode,
-          }
-        : replyParams;
-      const sent = await params.api.sendMessage(chatId, renderedText, sendParams);
-      const sentMessageId = sent?.message_id;
-      if (typeof sentMessageId !== "number" || !Number.isFinite(sentMessageId)) {
-        streamState.stopped = true;
-        params.warn?.("telegram stream preview stopped (missing message id from sendMessage)");
-        return false;
-      }
-      const normalizedMessageId = Math.trunc(sentMessageId);
-      if (sendGeneration !== generation) {
-        params.onSupersededPreview?.({
-          messageId: normalizedMessageId,
-          textSnapshot: renderedText,
-          parseMode: renderedParseMode,
-        });
-        return true;
-      }
-      streamMessageId = normalizedMessageId;
-      return true;
+      return sent;
     } catch (err) {
       streamState.stopped = true;
       params.warn?.(
@@ -166,6 +281,9 @@ export function createTelegramDraftStream(params: {
   const forceNewMessage = () => {
     generation += 1;
     streamMessageId = undefined;
+    if (usesDraftTransport) {
+      streamDraftId = allocateTelegramDraftId();
+    }
     lastSentText = "";
     lastSentParseMode = undefined;
     loop.resetPending();
@@ -178,6 +296,8 @@ export function createTelegramDraftStream(params: {
     update,
     flush: loop.flush,
     messageId: () => streamMessageId,
+    previewMode: () => (usesDraftTransport ? "draft" : "message"),
+    previewRevision: () => previewRevision,
     clear,
     stop,
     forceNewMessage,