From 19e794e34c55cbc5b719265a51664f3e18012ef4 Mon Sep 17 00:00:00 2001
From: Aidan <aidan@p0ntus.com>
Date: Thu, 3 Jul 2025 14:35:59 -0400
Subject: [PATCH] max queue size

---
 config/ai.ts                |  1 +
 src/commands/ai.ts          | 18 +++++++++++++++---
 src/locales/english.json    |  1 +
 src/locales/portuguese.json |  3 ++-
 4 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/config/ai.ts b/config/ai.ts
index 5f9a5cf..1cf4fdf 100644
--- a/config/ai.ts
+++ b/config/ai.ts
@@ -3,6 +3,7 @@ import type { ModelInfo } from "../src/commands/ai"
 export const defaultFlashModel = "gemma3:4b"
 export const defaultThinkingModel = "qwen3:4b"
 export const unloadModelAfterB = 12 // how many billion params until model is auto-unloaded
+export const maxUserQueueSize = 3
 
 export const models: ModelInfo[] = [
   {
diff --git a/src/commands/ai.ts b/src/commands/ai.ts
index 507f6a5..3acdc0b 100644
--- a/src/commands/ai.ts
+++ b/src/commands/ai.ts
@@ -41,7 +41,7 @@ import { ensureUserInDb } from "../utils/ensure-user"
 import * as schema from '../db/schema'
 import type { NodePgDatabase } from "drizzle-orm/node-postgres"
 import { eq, sql } from 'drizzle-orm'
-import { models, unloadModelAfterB } from "../../config/ai"
+import { models, unloadModelAfterB, maxUserQueueSize } from "../../config/ai"
 
 const spamwatchMiddleware = spamwatchMiddlewareModule(isOnSpamWatch)
 export const flash_model = process.env.flashModel || "gemma3:4b"
@@ -503,6 +503,7 @@ export default (bot: Telegraf<Context>, db: NodePgDatabase<typeof schema>) => {
     task: () => Promise<void>;
     ctx: TextContext;
     wasQueued: boolean;
+    userId: number;
   }
 
   const requestQueue: AiRequest[] = [];
@@ -567,6 +568,17 @@ export default (bot: Telegraf<Context>, db: NodePgDatabase<typeof schema>) => {
       return;
     }
 
+    const userId = ctx.from!.id;
+    const userQueueSize = requestQueue.filter(req => req.userId === userId).length;
+
+    if (userQueueSize >= maxUserQueueSize) {
+      await ctx.reply(Strings.ai.queueFull, {
+        parse_mode: 'Markdown',
+        ...(reply_to_message_id && { reply_parameters: { message_id: reply_to_message_id } })
+      });
+      return;
+    }
+
     const task = async () => {
       const modelLabel = getModelLabelByName(model);
       const replyGenerating = await ctx.reply(Strings.ai.askGenerating.replace("{model}", `\`${modelLabel}\``), {
@@ -578,14 +590,14 @@ export default (bot: Telegraf<Context>, db: NodePgDatabase<typeof schema>) => {
     };
 
     if (isProcessing) {
-      requestQueue.push({ task, ctx, wasQueued: true });
+      requestQueue.push({ task, ctx, wasQueued: true, userId: ctx.from!.id });
       const position = requestQueue.length;
       await ctx.reply(Strings.ai.inQueue.replace("{position}", String(position)), {
         parse_mode: 'Markdown',
         ...(reply_to_message_id && { reply_parameters: { message_id: reply_to_message_id } })
       });
     } else {
-      requestQueue.push({ task, ctx, wasQueued: false });
+      requestQueue.push({ task, ctx, wasQueued: false, userId: ctx.from!.id });
       processQueue();
     }
   }
diff --git a/src/locales/english.json b/src/locales/english.json
index dcb178d..b20e2ec 100644
--- a/src/locales/english.json
+++ b/src/locales/english.json
@@ -78,6 +78,7 @@
     "finishedThinking": "`🧠 Done thinking.`",
     "urlWarning": "\n\n⚠️ Note: The model cannot access or visit links!",
     "inQueue": "ℹ️ You are {position} in the queue.",
+    "queueFull": "🚫 You already have too many requests in the queue. Please wait for them to finish.",
     "startingProcessing": "✨ Starting to process your request...",
     "systemPrompt": "You are a friendly assistant called {botName}.\nCurrent Date/Time (UTC): {date}\n\n---\n\nUser message:\n{message}",
     "statusWaitingRender": "⏳ Streaming...",
diff --git a/src/locales/portuguese.json b/src/locales/portuguese.json
index 8697c2b..eda8cf4 100644
--- a/src/locales/portuguese.json
+++ b/src/locales/portuguese.json
@@ -76,6 +76,7 @@
         "finishedThinking": "`🧠 Pensamento concluido.`",
         "urlWarning": "\n\n⚠️ Nota: O modelo de IA não pode acessar ou visitar links!",
         "inQueue": "ℹ️ Você é o {position} na fila.",
+        "queueFull": "🚫 Você já tem muitas solicitações na fila. Por favor, espere que elas terminem.",
         "startingProcessing": "✨ Começando a processar o seu pedido...",
         "aiEnabled": "IA",
         "aiModel": "Modelo de IA",
@@ -191,4 +192,4 @@
         "requests": "*Total de requisições de IA:* {aiRequests}",
         "characters": "*Total de caracteres de IA:* {aiCharacters}\n_Isso é cerca de {bookCount} livros de texto!_"
     }
-}
+}
\ No newline at end of file