feat: persistent queue

bot.ts

@@ -1,83 +1,41 @@
-import {
+import { autoQuote, bold, Bot, Context, hydrateReply, ParseModeFlavor } from "./deps.ts";
-  autoQuote,
+import { fmt } from "./intl.ts";
-  autoRetry,
+import { getAllJobs, pushJob } from "./queue.ts";
-  bold,
+import { mySession, MySessionFlavor } from "./session.ts";
-  Bot,
-  Context,
-  DenoKVAdapter,
-  fmt,
-  hydrateReply,
-  ParseModeFlavor,
-  session,
-  SessionFlavor,
-} from "./deps.ts";
-import { fmtArray, formatOrdinal } from "./intl.ts";
-import { queue } from "./queue.ts";
-import { SdRequest } from "./sd.ts";
 
-type AppContext = ParseModeFlavor<Context> & SessionFlavor<SessionData>;
+export type MyContext = ParseModeFlavor<Context> & MySessionFlavor;
-
+export const bot = new Bot<MyContext>(Deno.env.get("TG_BOT_TOKEN") ?? "");
-interface SessionData {
-  global: {
-    adminUsernames: string[];
-    pausedReason: string | null;
-    sdApiUrl: string;
-    maxUserJobs: number;
-    maxJobs: number;
-    defaultParams?: Partial<SdRequest>;
-  };
-  user: {
-    steps: number;
-    detail: number;
-    batchSize: number;
-  };
-}
-
-export const bot = new Bot<AppContext>(Deno.env.get("TG_BOT_TOKEN") ?? "");
 bot.use(autoQuote);
 bot.use(hydrateReply);
-bot.api.config.use(autoRetry({ maxRetryAttempts: 5, maxDelaySeconds: 60 }));
+bot.use(mySession);
 
-const db = await Deno.openKv("./app.db");
+// Automatically retry bot requests if we get a 429 error
-
+bot.api.config.use(async (prev, method, payload, signal) => {
-const getDefaultGlobalSession = (): SessionData["global"] => ({
+  let remainingAttempts = 5;
-  adminUsernames: (Deno.env.get("ADMIN_USERNAMES") ?? "").split(",").filter(Boolean),
+  while (true) {
-  pausedReason: null,
+    const result = await prev(method, payload, signal);
-  sdApiUrl: Deno.env.get("SD_API_URL") ?? "http://127.0.0.1:7860/",
+    if (result.ok) return result;
-  maxUserJobs: 3,
+    if (result.error_code !== 429 || remainingAttempts <= 0) return result;
-  maxJobs: 20,
+    remainingAttempts -= 1;
-  defaultParams: {
+    const retryAfterMs = (result.parameters?.retry_after ?? 30) * 1000;
-    batch_size: 1,
+    await new Promise((resolve) => setTimeout(resolve, retryAfterMs));
-    n_iter: 1,
+  }
-    width: 128 * 2,
-    height: 128 * 3,
-    steps: 20,
-    cfg_scale: 9,
-    send_images: true,
-    negative_prompt: "boring_e621_fluffyrock_v4 boring_e621_v4",
-  },
 });
 
-bot.use(session<SessionData, AppContext>({
+// if error happened, try to reply to the user with the error
-  type: "multi",
+bot.use(async (ctx, next) => {
-  global: {
+  try {
-    getSessionKey: () => "global",
+    await next();
-    initial: getDefaultGlobalSession,
+  } catch (err) {
-    storage: new DenoKVAdapter(db),
+    try {
-  },
+      await ctx.reply(`Handling update failed: ${err}`, {
-  user: {
+        reply_to_message_id: ctx.message?.message_id,
-    initial: () => ({
+      });
-      steps: 20,
+    } catch {
-      detail: 8,
+      throw err;
-      batchSize: 2,
+    }
-    }),
+  }
-  },
+});
-}));
-
-export async function getGlobalSession(): Promise<SessionData["global"]> {
-  const entry = await db.get<SessionData["global"]>(["sessions", "global"]);
-  return entry.value ?? getDefaultGlobalSession();
-}
 
 bot.api.setMyShortDescription("I can generate furry images from text");
 bot.api.setMyDescription(
@@ -99,12 +57,13 @@ bot.command("txt2img", async (ctx) => {
   if (config.pausedReason != null) {
     return ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`);
   }
-  if (queue.length >= config.maxJobs) {
+  const jobs = await getAllJobs();
+  if (jobs.length >= config.maxJobs) {
     return ctx.reply(
       `The queue is full. Try again later. (Max queue size: ${config.maxJobs})`,
     );
   }
-  const jobCount = queue.filter((job) => job.userId === ctx.from.id).length;
+  const jobCount = jobs.filter((job) => job.user.id === ctx.from.id).length;
   if (jobCount >= config.maxUserJobs) {
     return ctx.reply(
       `You already have ${config.maxUserJobs} jobs in queue. Try again later.`,
@@ -113,36 +72,50 @@ bot.command("txt2img", async (ctx) => {
   if (!ctx.match) {
     return ctx.reply("Please describe what you want to see after the command");
   }
-  const place = queue.length + 1;
+  pushJob({
-  const queueMessage = await ctx.reply(`You are ${formatOrdinal(place)} in queue.`);
-  const userName = [ctx.from.first_name, ctx.from.last_name].filter(Boolean).join(" ");
-  const chatName = ctx.chat.type === "supergroup" || ctx.chat.type === "group"
-    ? ctx.chat.title
-    : "private chat";
-  queue.push({
     params: { prompt: ctx.match },
-    userId: ctx.from.id,
+    user: ctx.from,
-    userName,
+    chat: ctx.chat,
-    chatId: ctx.chat.id,
+    requestMessage: ctx.message,
-    chatName,
+    status: { type: "idle" },
-    requestMessageId: ctx.message.message_id,
-    statusMessageId: queueMessage.message_id,
   });
-  console.log(`Enqueued job for ${userName} in chat ${chatName}`);
+  console.log(
+    `Enqueued job ${jobs.length + 1} for ${ctx.from.first_name} in ${ctx.chat.type} chat:`,
+    ctx.match.replace(/\s+/g, " "),
+    "\n",
+  );
 });
 
-bot.command("queue", (ctx) => {
+bot.command("queue", async (ctx) => {
-  if (queue.length === 0) return ctx.reply("Queue is empty");
+  let jobs = await getAllJobs();
-  return ctx.replyFmt(
+  const getMessageText = () => {
-    fmt`Current queue:\n\n${
+    if (jobs.length === 0) return fmt`Queue is empty.`;
-      fmtArray(
+    const sortedJobs = [];
-        queue.map((job, index) =>
+    let place = 0;
-          fmt`${bold(index + 1)}. ${bold(job.userName)} in ${bold(job.chatName)}`
+    for (const job of jobs) {
-        ),
+      if (job.status.type === "idle") place += 1;
-        "\n",
+      sortedJobs.push({ ...job, place });
+    }
+    return fmt`Current queue:\n\n${
+      sortedJobs.map((job) =>
+        fmt`${job.place}. ${bold(job.user.first_name)} in ${job.chat.type} chat ${
+          job.status.type === "processing" ? `(${(job.status.progress * 100).toFixed(0)}%)` : ""
+        }\n`
       )
-    }`,
+    }`;
-  );
+  };
+  const message = await ctx.replyFmt(getMessageText());
+  handleFutureUpdates();
+  async function handleFutureUpdates() {
+    for (let idx = 0; idx < 12; idx++) {
+      await new Promise((resolve) => setTimeout(resolve, 5000));
+      jobs = await getAllJobs();
+      const formattedMessage = getMessageText();
+      await ctx.api.editMessageText(ctx.chat.id, message.message_id, formattedMessage.text, {
+        entities: formattedMessage.entities,
+      }).catch(() => undefined);
+    }
+  }
 });
 
 bot.command("pause", (ctx) => {
@@ -272,14 +245,17 @@ bot.command("setsdparam", (ctx) => {
 bot.command("sdparams", (ctx) => {
   if (!ctx.from?.username) return;
   const config = ctx.session.global;
-  return ctx.replyFmt(fmt`Current config:\n\n${
+  return ctx.replyFmt(
-    fmtArray(
+    fmt`Current config:\n\n${
       Object.entries(config.defaultParams ?? {}).map(([key, value]) =>
-        fmt`${bold(key)} = ${String(value)}`
+        fmt`${bold(key)} = ${String(value)}\n`
-      ),
-      "\n",
       )
-  }`);
+    }`,
+  );
+});
+
+bot.command("crash", () => {
+  throw new Error("Crash command used");
 });
 
 bot.catch((err) => {

deps.ts

@@ -2,5 +2,5 @@ export * from "https://deno.land/x/grammy@v1.18.1/mod.ts";
 export * from "https://deno.land/x/grammy_autoquote@v1.1.2/mod.ts";
 export * from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts";
 export * from "https://deno.land/x/grammy_storages@v2.3.1/denokv/src/mod.ts";
-export { autoRetry } from "https://esm.sh/@grammyjs/auto-retry@1.1.1";
+export * as types from "https://deno.land/x/grammy_types@v3.2.0/mod.ts";
-export * from "https://deno.land/x/zod/mod.ts";
+export * from "https://deno.land/x/ulid@v0.3.0/mod.ts";

intl.ts

@@ -8,26 +8,36 @@ export function formatOrdinal(n: number) {
   return `${n}th`;
 }
 
+type DeepArray<T> = Array<T | DeepArray<T>>;
+type StringLikes = DeepArray<FormattedString | string | number | null | undefined>;
+
 /**
- * Like `fmt` from `grammy_parse_mode` but accepts an array instead of template string.
+ * Like `fmt` from `grammy_parse_mode` but additionally accepts arrays.
  * @see https://deno.land/x/grammy_parse_mode@1.7.1/format.ts?source=#L182
  */
-export function fmtArray(
+export const fmt = (
-  stringLikes: FormattedString[],
+  rawStringParts: TemplateStringsArray | StringLikes,
-  separator = "",
+  ...stringLikes: StringLikes
-): FormattedString {
+): FormattedString => {
   let text = "";
-  const entities: ConstructorParameters<typeof FormattedString>[1] = [];
+  const entities: ConstructorParameters<typeof FormattedString>[1][] = [];
-  for (let i = 0; i < stringLikes.length; i++) {
+
-    const stringLike = stringLikes[i];
+  const length = Math.max(rawStringParts.length, stringLikes.length);
+  for (let i = 0; i < length; i++) {
+    for (let stringLike of [rawStringParts[i], stringLikes[i]]) {
+      if (Array.isArray(stringLike)) {
+        stringLike = fmt(stringLike);
+      }
+      if (stringLike instanceof FormattedString) {
         entities.push(
           ...stringLike.entities.map((e) => ({
             ...e,
             offset: e.offset + text.length,
           })),
         );
-    text += stringLike.toString();
+      }
-    if (i < stringLikes.length - 1) text += separator;
+      if (stringLike != null) text += stringLike.toString();
+    }
   }
   return new FormattedString(text, entities);
-}
+};

main.ts

@@ -1,8 +1,9 @@
 import "https://deno.land/std@0.201.0/dotenv/load.ts";
 import { bot } from "./bot.ts";
-import { processQueue } from "./queue.ts";
+import { processQueue, returnHangedJobs } from "./queue.ts";
 
 await Promise.all([
   bot.start(),
   processQueue(),
+  returnHangedJobs(),
 ]);

queue.ts

@@ -1,111 +1,167 @@
-import { InputFile, InputMediaBuilder } from "./deps.ts";
+import { InputFile, InputMediaBuilder, types } from "./deps.ts";
-import { bot, getGlobalSession } from "./bot.ts";
+import { bot } from "./bot.ts";
+import { getGlobalSession } from "./session.ts";
 import { formatOrdinal } from "./intl.ts";
-import { SdProgressResponse, SdRequest, txt2img } from "./sd.ts";
+import { SdRequest, txt2img } from "./sd.ts";
 import { extFromMimeType, mimeTypeFromBase64 } from "./mimeType.ts";
-
+import { Model, Store } from "./store.ts";
-export const queue: Job[] = [];
 
 interface Job {
   params: Partial<SdRequest>;
-  userId: number;
+  user: types.User;
-  userName: string;
+  chat: types.Chat.PrivateChat | types.Chat.GroupChat | types.Chat.SupergroupChat;
-  chatId: number;
+  requestMessage: types.Message & types.Message.TextMessage;
-  chatName: string;
+  statusMessage?: types.Message & types.Message.TextMessage;
-  requestMessageId: number;
+  status: { type: "idle" } | { type: "processing"; progress: number; updatedDate: Date };
-  statusMessageId: number;
+}
+
+const db = await Deno.openKv("./app.db");
+
+const jobStore = new Store<Job>(db, "job");
+
+export async function pushJob(job: Job) {
+  await jobStore.create(job);
+}
+
+async function takeJob(): Promise<Model<Job> | null> {
+  const jobs = await jobStore.list();
+  const job = jobs.find((job) => job.value.status.type === "idle");
+  if (!job) return null;
+  await job.update({ status: { type: "processing", progress: 0, updatedDate: new Date() } });
+  return job;
+}
+
+export async function getAllJobs(): Promise<Array<Job>> {
+  return await jobStore.list().then((jobs) => jobs.map((job) => job.value));
 }
 
 export async function processQueue() {
   while (true) {
-    const job = queue.shift();
+    const job = await takeJob();
     if (!job) {
       await new Promise((resolve) => setTimeout(resolve, 1000));
       continue;
     }
-    for (const [index, job] of queue.entries()) {
+    let place = 0;
-      const place = index + 1;
+    for (const job of await jobStore.list()) {
-      await bot.api
+      if (job.value.status.type === "idle") place += 1;
-        .editMessageText(
+      if (place === 0) continue;
-          job.chatId,
+      const statusMessageText = `You are ${formatOrdinal(place)} in queue.`;
-          job.statusMessageId,
+      if (!job.value.statusMessage) {
-          `You are ${formatOrdinal(place)} in queue.`,
+        await bot.api.sendMessage(job.value.chat.id, statusMessageText, {
+          reply_to_message_id: job.value.requestMessage.message_id,
+        }).catch(() => undefined)
+          .then((message) => job.update({ statusMessage: message }));
+      } else {
+        await bot.api.editMessageText(
+          job.value.chat.id,
+          job.value.statusMessage.message_id,
+          statusMessageText,
         )
-        .catch(() => {});
+          .catch(() => undefined);
+      }
     }
     try {
+      if (job.value.statusMessage) {
         await bot.api
-        .deleteMessage(job.chatId, job.statusMessageId)
+          .deleteMessage(job.value.chat.id, job.value.statusMessage?.message_id)
-        .catch(() => {});
+          .catch(() => undefined)
-      const progressMessage = await bot.api.sendMessage(
+          .then(() => job.update({ statusMessage: undefined }));
-        job.chatId,
+      }
+      await bot.api.sendMessage(
+        job.value.chat.id,
         "Generating your prompt now...",
-        { reply_to_message_id: job.requestMessageId },
+        { reply_to_message_id: job.value.requestMessage.message_id },
-      );
+      ).then((message) => job.update({ statusMessage: message }));
-      const onProgress = (progress: SdProgressResponse) => {
+      const config = await getGlobalSession();
+      const response = await txt2img(
+        config.sdApiUrl,
+        { ...config.defaultParams, ...job.value.params },
+        (progress) => {
+          job.update({
+            status: { type: "processing", progress: progress.progress, updatedDate: new Date() },
+          });
+          if (job.value.statusMessage) {
             bot.api
               .editMessageText(
-            job.chatId,
+                job.value.chat.id,
-            progressMessage.message_id,
+                job.value.statusMessage.message_id,
                 `Generating your prompt now... ${
                   Math.round(
                     progress.progress * 100,
                   )
                 }%`,
               )
-          .catch(() => {});
+              .catch(() => undefined);
-      };
+          }
-      const config = await getGlobalSession();
+        },
-      const response = await txt2img(
-        config.sdApiUrl,
-        { ...config.defaultParams, ...job.params },
-        onProgress,
       );
       console.log(
-        `Generated ${response.images.length} images (${
+        `Finished job for ${job.value.user.first_name} in ${job.value.chat.type} chat`,
-          response.images
-            .map((image) => (image.length / 1024).toFixed(0) + "kB")
-            .join(", ")
-        }) for ${job.userName} in ${job.chatName}: ${job.params.prompt?.replace(/\s+/g, " ")}`,
       );
+      if (job.value.statusMessage) {
         await bot.api.editMessageText(
-        job.chatId,
+          job.value.chat.id,
-        progressMessage.message_id,
+          job.value.statusMessage.message_id,
           `Uploading your images...`,
-      ).catch(() => {});
+        ).catch(() => undefined);
+      }
       const inputFiles = await Promise.all(
         response.images.map(async (imageBase64, idx) => {
           const mimeType = mimeTypeFromBase64(imageBase64);
-          const imageBlob = await fetch(`data:${mimeType};base64,${imageBase64}`).then((resp) =>
+          const imageBlob = await fetch(`data:${mimeType};base64,${imageBase64}`)
-            resp.blob()
+            .then((resp) => resp.blob());
-          );
-          console.log(
-            `Uploading image ${idx + 1} of ${response.images.length} (${
-              (imageBlob.size / 1024).toFixed(0)
-            }kB)`,
-          );
           return InputMediaBuilder.photo(
-            new InputFile(imageBlob, `${idx}.${extFromMimeType(mimeType)}`),
+            new InputFile(imageBlob, `image_${idx}.${extFromMimeType(mimeType)}`),
           );
         }),
       );
-      await bot.api.sendMediaGroup(job.chatId, inputFiles, {
+      if (job.value.statusMessage) {
-        reply_to_message_id: job.requestMessageId,
-      });
         await bot.api
-        .deleteMessage(job.chatId, progressMessage.message_id)
+          .deleteMessage(job.value.chat.id, job.value.statusMessage.message_id)
-        .catch(() => {});
+          .catch(() => undefined).then(() => job.update({ statusMessage: undefined }));
-      console.log(`${queue.length} jobs remaining`);
+      }
+      await bot.api.sendMediaGroup(job.value.chat.id, inputFiles, {
+        reply_to_message_id: job.value.requestMessage.message_id,
+      });
+      await job.delete();
     } catch (err) {
       console.error(
-        `Failed to generate image for ${job.userName} in ${job.chatName}: ${job.params.prompt} - ${err}`,
+        `Failed to generate an image for ${job.value.user.first_name} in ${job.value.chat.type} chat: ${err}`,
       );
-      await bot.api
+      const errorMessage = await bot.api
-        .sendMessage(job.chatId, err.toString(), {
+        .sendMessage(job.value.chat.id, err.toString(), {
-          reply_to_message_id: job.requestMessageId,
+          reply_to_message_id: job.value.requestMessage.message_id,
         })
-        .catch(() => bot.api.sendMessage(job.chatId, err.toString()))
+        .catch(() => undefined);
-        .catch(() => {});
+      if (errorMessage) {
+        if (job.value.statusMessage) {
+          await bot.api
+            .deleteMessage(job.value.chat.id, job.value.statusMessage.message_id)
+            .catch(() => undefined)
+            .then(() => job.update({ statusMessage: undefined }));
+        }
+        job.update({ status: { type: "idle" } });
+      } else {
+        await job.delete();
+      }
+    }
+  }
+}
+
+export async function returnHangedJobs() {
+  while (true) {
+    await new Promise((resolve) => setTimeout(resolve, 5000));
+    const jobs = await jobStore.list();
+    for (const job of jobs) {
+      if (job.value.status.type === "idle") continue;
+      // if job wasn't updated for 1 minute, return it to the queue
+      if (job.value.status.updatedDate.getTime() < Date.now() - 60 * 1000) {
+        console.log(
+          `Returning hanged job for ${job.value.user.first_name} in ${job.value.chat.type} chat`,
+        );
+        await job.update({ status: { type: "idle" } });
+      }
     }
   }
 }

session.ts

@@ -0,0 +1,78 @@
+import { Context, DenoKVAdapter, session, SessionFlavor } from "./deps.ts";
+import { SdRequest } from "./sd.ts";
+
+export type MySessionFlavor = SessionFlavor<SessionData>;
+
+export interface SessionData {
+  global: GlobalData;
+  chat: ChatData;
+  user: UserData;
+}
+
+export interface GlobalData {
+  adminUsernames: string[];
+  pausedReason: string | null;
+  sdApiUrl: string;
+  maxUserJobs: number;
+  maxJobs: number;
+  defaultParams?: Partial<SdRequest>;
+}
+
+export interface ChatData {
+  language: string;
+}
+
+export interface UserData {
+  steps: number;
+  detail: number;
+  batchSize: number;
+}
+
+const globalDb = await Deno.openKv("./app.db");
+
+const globalDbAdapter = new DenoKVAdapter<GlobalData>(globalDb);
+
+const getDefaultGlobalData = (): GlobalData => ({
+  adminUsernames: (Deno.env.get("ADMIN_USERNAMES") ?? "").split(",").filter(Boolean),
+  pausedReason: null,
+  sdApiUrl: Deno.env.get("SD_API_URL") ?? "http://127.0.0.1:7860/",
+  maxUserJobs: 3,
+  maxJobs: 20,
+  defaultParams: {
+    batch_size: 1,
+    n_iter: 1,
+    width: 128 * 2,
+    height: 128 * 3,
+    steps: 20,
+    cfg_scale: 9,
+    send_images: true,
+    negative_prompt: "boring_e621_fluffyrock_v4 boring_e621_v4",
+  },
+});
+
+export const mySession = session<SessionData, Context & MySessionFlavor>({
+  type: "multi",
+  global: {
+    getSessionKey: () => "global",
+    initial: getDefaultGlobalData,
+    storage: globalDbAdapter,
+  },
+  chat: {
+    initial: () => ({
+      language: "en",
+    }),
+  },
+  user: {
+    getSessionKey: (ctx) => ctx.from?.id.toFixed(),
+    initial: () => ({
+      steps: 20,
+      detail: 8,
+      batchSize: 2,
+    }),
+  },
+});
+
+export async function getGlobalSession(): Promise<GlobalData> {
+  const data = await globalDbAdapter.read("global");
+  return data ?? getDefaultGlobalData();
+}

store.ts

@@ -0,0 +1,76 @@
+import { ulid } from "./deps.ts";
+
+export class Store<T extends object> {
+  constructor(
+    private readonly db: Deno.Kv,
+    private readonly storeKey: Deno.KvKeyPart,
+  ) {
+  }
+
+  async create(value: T): Promise<Model<T>> {
+    const id = ulid();
+    await this.db.set([this.storeKey, id], value);
+    return new Model(this.db, this.storeKey, id, value);
+  }
+
+  async get(id: Deno.KvKeyPart): Promise<Model<T> | null> {
+    const entry = await this.db.get<T>([this.storeKey, id]);
+    if (entry.versionstamp == null) return null;
+    return new Model(this.db, this.storeKey, id, entry.value);
+  }
+
+  async list(): Promise<Array<Model<T>>> {
+    const models: Array<Model<T>> = [];
+    for await (const entry of this.db.list<T>({ prefix: [this.storeKey] })) {
+      models.push(new Model(this.db, this.storeKey, entry.key[1], entry.value));
+    }
+    return models;
+  }
+}
+
+export class Model<T extends object> {
+  #value: T;
+
+  constructor(
+    private readonly db: Deno.Kv,
+    private readonly storeKey: Deno.KvKeyPart,
+    private readonly entryKey: Deno.KvKeyPart,
+    value: T,
+  ) {
+    this.#value = value;
+  }
+
+  get value(): T {
+    return this.#value;
+  }
+
+  async get(): Promise<T | null> {
+    const entry = await this.db.get<T>([this.storeKey, this.entryKey]);
+    if (entry.versionstamp == null) return null;
+    this.#value = entry.value;
+    return entry.value;
+  }
+
+  async set(value: T): Promise<T> {
+    await this.db.set([this.storeKey, this.entryKey], value);
+    this.#value = value;
+    return value;
+  }
+
+  async update(value: Partial<T> | ((value: T) => T)): Promise<T | null> {
+    const entry = await this.db.get<T>([this.storeKey, this.entryKey]);
+    if (entry.versionstamp == null) return null;
+    if (typeof value === "function") {
+      entry.value = value(entry.value);
+    } else {
+      entry.value = { ...entry.value, ...value };
+    }
+    await this.db.set([this.storeKey, this.entryKey], entry.value);
+    this.#value = entry.value;
+    return entry.value;
+  }
+
+  async delete(): Promise<void> {
+    await this.db.delete([this.storeKey, this.entryKey]);
+  }
+}