From 6b0d7b7198f4ffc31aa77b91c3668e7536177c24 Mon Sep 17 00:00:00 2001
From: pinks <lisq@cock.li>
Date: Mon, 11 Sep 2023 19:07:46 +0200
Subject: [PATCH] better tg api retry logic

---
 bot/mod.ts                   | 16 ++++++++++++++--
 tasks/processJobs.ts         |  5 ++++-
 tasks/updateJobStatusMsgs.ts |  1 +
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/bot/mod.ts b/bot/mod.ts
index ecd33ce..deb234c 100644
--- a/bot/mod.ts
+++ b/bot/mod.ts
@@ -6,8 +6,16 @@ import { txt2imgCommand, txt2imgQuestion } from "./txt2imgCommand.ts";
 
 export const logger = () => Log.getLogger();
 
+type WithRetryApi<T extends Grammy.RawApi> = {
+  [M in keyof T]: T[M] extends (args: infer P, ...rest: infer A) => infer R
+    ? (args: P extends object ? P & { maxAttempts?: number } : P, ...rest: A) => R
+    : T[M];
+};
+
 export type Context = GrammyParseMode.ParseModeFlavor<Grammy.Context> & SessionFlavor;
-export const bot = new Grammy.Bot<Context>(Deno.env.get("TG_BOT_TOKEN") ?? "");
+export const bot = new Grammy.Bot<Context, Grammy.Api<WithRetryApi<Grammy.RawApi>>>(
+  Deno.env.get("TG_BOT_TOKEN") ?? "",
+);
 bot.use(GrammyAutoQuote.autoQuote);
 bot.use(GrammyParseMode.hydrateReply);
 bot.use(session);
@@ -18,6 +26,7 @@ bot.catch((err) => {
 
 // Automatically retry bot requests if we get a "too many requests" or telegram internal error
 bot.api.config.use(async (prev, method, payload, signal) => {
+  const maxAttempts = payload && ("maxAttempts" in payload) ? payload.maxAttempts ?? 3 : 3;
   let attempt = 0;
   while (true) {
     attempt++;
@@ -25,10 +34,13 @@ bot.api.config.use(async (prev, method, payload, signal) => {
     if (
       result.ok ||
       ![429, 500].includes(result.error_code) ||
-      attempt >= 5
+      attempt >= maxAttempts
     ) {
       return result;
     }
+    logger().warning(
+      `Retrying ${method} after attempt ${attempt} failed with ${result.error_code} error`,
+    );
     const retryAfterMs = (result.parameters?.retry_after ?? (attempt * 5)) * 1000;
     await new Promise((resolve) => setTimeout(resolve, retryAfterMs));
   }
diff --git a/tasks/processJobs.ts b/tasks/processJobs.ts
index 5df61bb..1cfff8f 100644
--- a/tasks/processJobs.ts
+++ b/tasks/processJobs.ts
@@ -121,6 +121,7 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
           `Generating your prompt now... ${
             (progress.progress * 100).toFixed(0)
           }% using ${worker.name}`,
+          { maxAttempts: 1 },
         ).catch(() => undefined);
       }
       await job.update({
@@ -130,7 +131,7 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
           worker: worker.name,
           updatedDate: new Date(),
         },
-      }).catch(() => undefined);
+      }, { maxAttempts: 1 }).catch(() => undefined);
     },
   );
 
@@ -140,6 +141,7 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
       job.value.reply.chat.id,
       job.value.reply.message_id,
       `Uploading your images...`,
+      { maxAttempts: 1 },
     ).catch(() => undefined);
   }
 
@@ -185,6 +187,7 @@ async function processJob(job: IKV.Model<JobSchema>, worker: WorkerData, config:
   // send the result to telegram
   const resultMessage = await bot.api.sendMediaGroup(job.value.request.chat.id, inputFiles, {
     reply_to_message_id: job.value.request.message_id,
+    maxAttempts: 5,
   });
   // send caption in separate message if it couldn't fit
   if (caption.text.length > 1024 && caption.text.length <= 4096) {
diff --git a/tasks/updateJobStatusMsgs.ts b/tasks/updateJobStatusMsgs.ts
index 2fa6149..a2a89c2 100644
--- a/tasks/updateJobStatusMsgs.ts
+++ b/tasks/updateJobStatusMsgs.ts
@@ -19,6 +19,7 @@ export async function updateJobStatusMsgs(): Promise<never> {
           job.value.reply.chat.id,
           job.value.reply.message_id,
           `You are ${formatOrdinal(index + 1)} in queue.`,
+          { maxAttempts: 1 },
         ).catch(() => undefined);
       }
     } catch (err) {