From 86bdfec6b70808ed2fdb3c609da534f06cd65672 Mon Sep 17 00:00:00 2001
From: Ayaan Zaidi <hi@obviy.us>
Date: Tue, 30 Jun 2026 10:17:08 -0700
Subject: [PATCH] refactor(telegram): simplify album prompt selection

---
 .../telegram/src/bot-handlers.runtime.ts      | 56 ++++++-------------
 .../telegram/src/bot-message-context.types.ts |  2 -
 2 files changed, 18 insertions(+), 40 deletions(-)

diff --git a/extensions/telegram/src/bot-handlers.runtime.ts b/extensions/telegram/src/bot-handlers.runtime.ts
index f30623e286d..135a39d47e5 100644
--- a/extensions/telegram/src/bot-handlers.runtime.ts
+++ b/extensions/telegram/src/bot-handlers.runtime.ts
@@ -245,6 +245,7 @@ export const registerTelegramHandlers = ({
     dispatchDedupeKeys: string[];
     spooledReplayParticipants: TelegramSpooledReplayDeferredParticipant[];
   };
+  type PromptContextMessageSelection = ReadonlyMap<string, "include" | "exclude">;
 
   const mediaGroupBuffer = new Map<string, BufferedMediaGroupEntry>();
   const mediaGroupProcessingByKey = new Map<string, Promise<void>>();
@@ -950,8 +951,10 @@ export const registerTelegramHandlers = ({
       }
 
       const allMedia: TelegramMediaRef[] = [];
+      const promptContextMessageSelection = new Map<string, "include" | "exclude">();
       let skippedCount = 0;
       for (const { ctx, msg } of entry.messages) {
+        const sourceMessageId = String(msg.message_id);
         let media;
         try {
           media = await resolveMedia({
@@ -966,6 +969,7 @@ export const registerTelegramHandlers = ({
           runtime.log?.(
             warn(`media group: skipping photo that failed to fetch: ${String(mediaErr)}`),
           );
+          promptContextMessageSelection.set(sourceMessageId, "exclude");
           skippedCount++;
           continue;
         }
@@ -974,11 +978,11 @@ export const registerTelegramHandlers = ({
             path: media.path,
             contentType: media.contentType,
             stickerMetadata: media.stickerMetadata,
-            ...(typeof msg.message_id === "number"
-              ? { sourceMessageId: String(msg.message_id) }
-              : {}),
+            sourceMessageId,
           });
+          promptContextMessageSelection.set(sourceMessageId, "include");
         } else {
+          promptContextMessageSelection.set(sourceMessageId, "exclude");
           skippedCount++;
         }
       }
@@ -1003,25 +1007,11 @@ export const registerTelegramHandlers = ({
         }).catch(() => {});
       }
 
-      // Only successfully downloaded album items have hydrated prompt media.
-      // Skipped siblings still have raw telegram:file refs in the cache.
-      const allAlbumMessageIds = entry.messages.flatMap(({ msg }) =>
-        typeof msg.message_id === "number" ? [String(msg.message_id)] : [],
-      );
-      const hydratedAlbumMessageIds = allMedia.flatMap((media) =>
-        media.sourceMessageId ? [media.sourceMessageId] : [],
-      );
-      const hydratedAlbumMessageIdSet = new Set(hydratedAlbumMessageIds);
-      const skippedAlbumMessageIds = allAlbumMessageIds.filter(
-        (messageId) => !hydratedAlbumMessageIdSet.has(messageId),
-      );
-
       const result = await processMessageWithReplyChain({
         ctx: primaryEntry.ctx,
         msg: primaryEntry.msg,
         allMedia,
-        promptContextExtraMessageIds: hydratedAlbumMessageIds,
-        promptContextExcludedMessageIds: skippedAlbumMessageIds,
+        promptContextMessageSelection,
         storeAllowFrom: entry.storeAllowFrom,
         options: {
           ...promptContextBoundaryOptions(entry.promptContextMinTimestampMs),
@@ -1249,8 +1239,7 @@ export const registerTelegramHandlers = ({
     replyChainNodes: TelegramCachedMessageNode[],
     options?: TelegramMessageContextOptions,
     mediaByMessageId?: ReadonlyMap<string, TelegramMediaRef>,
-    extraMessageIds?: readonly string[],
-    excludedMessageIds?: readonly string[],
+    selectedMessageIds?: PromptContextMessageSelection,
   ): Promise<TelegramPromptContextEntry[]> => {
     const isGroup = msg.chat.type === "group" || msg.chat.type === "supergroup";
     const groupHistoryContextMode = isGroup
@@ -1312,22 +1301,18 @@ export const registerTelegramHandlers = ({
         entry.node.messageId ? [[entry.node.messageId, entry] as const] : [],
       ),
     );
-    const excludedMessageIdSet = new Set(excludedMessageIds ?? []);
-    for (const excludedMessageId of excludedMessageIdSet) {
-      conversationContextById.delete(excludedMessageId);
-    }
-    for (const extraMessageId of new Set(extraMessageIds ?? [])) {
-      if (
-        extraMessageId === messageId ||
-        excludedMessageIdSet.has(extraMessageId) ||
-        conversationContextById.has(extraMessageId)
-      ) {
+    for (const [selectedMessageId, selection] of selectedMessageIds ?? []) {
+      if (selection === "exclude") {
+        conversationContextById.delete(selectedMessageId);
+        continue;
+      }
+      if (selectedMessageId === messageId || conversationContextById.has(selectedMessageId)) {
         continue;
       }
       const node = await messageCache.get({
         accountId,
         chatId: msg.chat.id,
-        messageId: extraMessageId,
+        messageId: selectedMessageId,
       });
       if (node?.messageId) {
         conversationContextById.set(node.messageId, { node });
@@ -1419,8 +1404,7 @@ export const registerTelegramHandlers = ({
     ctx: TelegramContext;
     msg: Message;
     allMedia: TelegramMediaRef[];
-    promptContextExtraMessageIds?: string[];
-    promptContextExcludedMessageIds?: string[];
+    promptContextMessageSelection?: PromptContextMessageSelection;
     storeAllowFrom: string[];
     options?: TelegramMessageContextOptions;
     dispatchDedupeKeys?: string[];
@@ -1525,8 +1509,7 @@ export const registerTelegramHandlers = ({
         replyChainNodes,
         params.options,
         promptContextMediaByMessageId,
-        params.promptContextExtraMessageIds,
-        params.promptContextExcludedMessageIds,
+        params.promptContextMessageSelection,
       );
       const result = await processMessage(
         params.ctx,
@@ -2318,9 +2301,6 @@ export const registerTelegramHandlers = ({
             path: media.path,
             contentType: media.contentType,
             stickerMetadata: media.stickerMetadata,
-            ...(typeof msg.message_id === "number"
-              ? { sourceMessageId: String(msg.message_id) }
-              : {}),
           },
         ]
       : [];
diff --git a/extensions/telegram/src/bot-message-context.types.ts b/extensions/telegram/src/bot-message-context.types.ts
index ba0a2adad2c..d917f53a873 100644
--- a/extensions/telegram/src/bot-message-context.types.ts
+++ b/extensions/telegram/src/bot-message-context.types.ts
@@ -16,8 +16,6 @@ export type TelegramMediaRef = {
   path: string;
   contentType?: string;
   stickerMetadata?: StickerMetadata;
-  // Album flushes process one captioned primary message while sibling messages
-  // remain in the prompt cache; keep their ids so context can hydrate each one.
   sourceMessageId?: string;
 };