small changes

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));
   }
@@ -62,7 +74,7 @@ bot.api.setMyCommands([
 bot.command("start", (ctx) => ctx.reply("Hello! Use the /txt2img command to generate an image"));
 
 bot.command("txt2img", txt2imgCommand);
-bot.use(txt2imgQuestion.middleware() as any);
+bot.use(txt2imgQuestion.middleware());
 
 bot.command("queue", queueCommand);
 

bot/queueCommand.ts

@@ -15,7 +15,6 @@ export async function queueCommand(ctx: Grammy.CommandContext<Context>) {
     const waitingJobs = await jobStore.getBy("status.type", "waiting")
       .then((jobs) => jobs.map((job, index) => ({ ...job.value, place: index + 1 })));
     const jobs = [...processingJobs, ...waitingJobs];
-    const config = ctx.session.global;
     const { bold } = GrammyParseMode;
     return fmt([
       "Current queue:\n",
@@ -40,7 +39,7 @@ export async function queueCommand(ctx: Grammy.CommandContext<Context>) {
         ])
         : ["Queue is empty.\n"],
       "\nActive workers:\n",
-      ...config.workers.flatMap((worker) => [
+      ...ctx.session.global.workers.flatMap((worker) => [
         runningWorkers.has(worker.name) ? "✅ " : "☠️ ",
         fmt`${bold(worker.name)} `,
         `(max ${(worker.maxResolution / 1000000).toFixed(1)} Mpx) `,

bot/txt2imgCommand.ts

@@ -4,11 +4,11 @@ import { jobStore } from "../db/jobStore.ts";
 import { parsePngInfo } from "../sd.ts";
 import { Context, logger } from "./mod.ts";
 
-export const txt2imgQuestion = new GrammyStatelessQ.StatelessQuestion(
+export const txt2imgQuestion = new GrammyStatelessQ.StatelessQuestion<Context>(
   "txt2img",
   async (ctx) => {
     if (!ctx.message.text) return;
-    await txt2img(ctx as any, ctx.message.text, false);
+    await txt2img(ctx, ctx.message.text, false);
   },
 );
 
@@ -22,24 +22,23 @@ async function txt2img(ctx: Context, match: string, includeRepliedTo: boolean):
     return;
   }
 
-  const config = ctx.session.global;
-  if (config.pausedReason != null) {
-    await ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`);
+  if (ctx.session.global.pausedReason != null) {
+    await ctx.reply(`I'm paused: ${ctx.session.global.pausedReason || "No reason given"}`);
     return;
   }
 
   const jobs = await jobStore.getBy("status.type", "waiting");
-  if (jobs.length >= config.maxJobs) {
+  if (jobs.length >= ctx.session.global.maxJobs) {
     await ctx.reply(
-      `The queue is full. Try again later. (Max queue size: ${config.maxJobs})`,
+      `The queue is full. Try again later. (Max queue size: ${ctx.session.global.maxJobs})`,
     );
     return;
   }
 
   const userJobs = jobs.filter((job) => job.value.request.from.id === ctx.message?.from?.id);
-  if (userJobs.length >= config.maxUserJobs) {
+  if (userJobs.length >= ctx.session.global.maxUserJobs) {
     await ctx.reply(
-      `You already have ${config.maxUserJobs} jobs in queue. Try again later.`,
+      `You already have ${ctx.session.global.maxUserJobs} jobs in queue. Try again later.`,
     );
     return;
   }
@@ -48,6 +47,7 @@ async function txt2img(ctx: Context, match: string, includeRepliedTo: boolean):
   const repliedToMsg = ctx.message.reply_to_message;
   const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
   if (includeRepliedTo && repliedToText) {
+    // TODO: remove bot command from replied to text
     const originalParams = parsePngInfo(repliedToText);
     params = {
       ...originalParams,

deps.ts

@@ -11,7 +11,7 @@ export * as GrammyTypes from "https://deno.land/x/grammy_types@v3.2.0/mod.ts";
 export * as GrammyAutoQuote from "https://deno.land/x/grammy_autoquote@v1.1.2/mod.ts";
 export * as GrammyParseMode from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts";
 export * as GrammyKvStorage from "https://deno.land/x/grammy_storages@v2.3.1/denokv/src/mod.ts";
-export * as GrammyStatelessQ from "npm:@grammyjs/stateless-question";
+export * as GrammyStatelessQ from "https://deno.land/x/grammy_stateless_question_alpha@v3.0.3/mod.ts";
 export * as FileType from "npm:file-type@18.5.0";
 // @deno-types="./types/png-chunks-extract.d.ts"
 export * as PngChunksExtract from "npm:png-chunks-extract@1.0.0";

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) {

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) {

utils.ts

@@ -105,7 +105,7 @@ const languageToFlagMap: Record<string, string> = {
   "lb": "🇱🇺", // Luxembourgish - Luxembourg
 };
 
-export function getFlagEmoji(countryCode?: string): string | undefined {
-  if (!countryCode) return;
-  return languageToFlagMap[countryCode];
+export function getFlagEmoji(languageCode?: string): string | undefined {
+  if (!languageCode) return;
+  return languageToFlagMap[languageCode];
 }