diff --git a/README.md b/README.md index c381d9b..d610427 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ Telegram-бот для генерации и редактирования изо - **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI. - **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов. - **Оплата**: 5 Telegram Stars за обычную очередь, 15 за премиум; админам бесплатно, возврат при ошибке. -- **Премиум-очередь**: выполняется раньше обычной; при пустой премиум-очереди работает обычная. +- **Премиум-очередь**: чередуется с обычной (одна премиум, одна обычная); если одна пуста — работает другая. - Глобальная очередь с одной активной задачей на пользователя. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI. - SQLite: запросы, статусы, seed, пути к файлам. @@ -135,7 +135,7 @@ PREMIUM_STARS_PRICE=15 - Перед оплатой бот показывает количество людей в обычной и премиум-очередях. - Пользователь выбирает очередь кнопкой, после чего получает счёт. -- Премиум-задачи выполняются раньше обычных. +- Премиум- и обычные задачи чередуются: одна премиум, одна обычная; если одна очередь пуста, берётся из другой. - Админы из `ADMIN_USER_IDS` генерируют бесплатно и попадают в премиум-очередь. - Не-админу выставляется счёт через Telegram Stars (`XTR`). - Фильтрация промпта выполняется до выставления счёта. diff --git a/checklist.md b/checklist.md index e9dc6f6..1d8caaa 100644 --- a/checklist.md +++ b/checklist.md @@ -8,7 +8,7 @@ - Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15) и оплата Stars (16). - Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история. -- 48 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, вспомогательные функции ComfyUI. +- 49 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, вспомогательные функции ComfyUI. - Не проверено: реальная оплата Stars в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. ## 0. Допущения и решения по умолчанию @@ -332,7 +332,7 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); ## 16. Оплата Stars - [x] `PAYMENTS_ENABLED` / `STARS_PRICE`, админы бесплатно. -- [x] Премиум-очередь: `PREMIUM_STARS_PRICE`, приоритетное выполнение, счётчики очередей до оплаты. +- [x] Премиум-очередь: `PREMIUM_STARS_PRICE`, чередование с обычной, счётчики очередей до оплаты. - [x] Invoice `XTR` + `pre_checkout_query` + `successful_payment`. - [x] `pending_payments`: payload, сумма, статус, charge_id, request_id. - [x] Возврат звёзд при ошибке постановки в очередь и генерации. diff --git a/queue_manager.py b/queue_manager.py index cf00884..79dbe52 100644 --- a/queue_manager.py +++ b/queue_manager.py @@ -63,6 +63,7 @@ class QueueManager: self._premium: deque[int] = deque() self._normal: deque[int] = deque() + self._next_premium = True self._active_users: set[int] = set() self._wake = asyncio.Event() self._lock = asyncio.Lock() @@ -159,7 +160,43 @@ class QueueManager: return await self._db.queue_counts() async def position(self, request_id: int) -> int: - return await self._db.count_ahead(request_id) + request = await self._db.get_request(request_id) + if request is not None and request.status == "running": + return 0 + + counts = await self._db.queue_counts() + running = max( + 0, + sum(counts.values()) - len(self._premium) - len(self._normal), + ) + premium = list(self._premium) + normal = list(self._normal) + prefer_premium = self._next_premium + premium_index = 0 + normal_index = 0 + ahead = running + while premium_index < len(premium) or normal_index < len(normal): + if prefer_premium and premium_index < len(premium): + current = premium[premium_index] + premium_index += 1 + taken_premium = True + elif not prefer_premium and normal_index < len(normal): + current = normal[normal_index] + normal_index += 1 + taken_premium = False + elif premium_index < len(premium): + current = premium[premium_index] + premium_index += 1 + taken_premium = True + else: + current = normal[normal_index] + normal_index += 1 + taken_premium = False + if current == request_id: + return ahead + ahead += 1 + prefer_premium = not taken_premium + return ahead async def _worker(self, index: int) -> None: logger.info("Воркер %s запущен", index) @@ -174,14 +211,25 @@ class QueueManager: async def _next_request(self) -> int | None: while not self._stopping: - if self._premium: - return self._premium.popleft() - if self._normal: - return self._normal.popleft() + if self._next_premium: + if self._premium: + return self._take(self._premium, premium=True) + if self._normal: + return self._take(self._normal, premium=False) + else: + if self._normal: + return self._take(self._normal, premium=False) + if self._premium: + return self._take(self._premium, premium=True) self._wake.clear() await self._wake.wait() return None + def _take(self, queue: deque[int], *, premium: bool) -> int: + request_id = queue.popleft() + self._next_premium = not premium + return request_id + async def _process(self, request_id: int) -> None: request = await self._db.get_request(request_id) if request is None: diff --git a/tests/test_queue.py b/tests/test_queue.py index b7d2fe6..1a53984 100644 --- a/tests/test_queue.py +++ b/tests/test_queue.py @@ -270,3 +270,27 @@ async def test_premium_queue_goes_first(db: Database, tmp_path: Path) -> None: assert second == normal.id await queue.stop() + + +async def test_premium_and_normal_alternate(db: Database, tmp_path: Path) -> None: + events: list[QueueEvent] = [] + queue = make_queue(db, make_settings(tmp_path), events) + + normal_1, _ = await queue.enqueue(user_id=1, mode="t2i", prompt="n1") + normal_2, _ = await queue.enqueue(user_id=2, mode="t2i", prompt="n2") + premium_1, _ = await queue.enqueue( + user_id=3, mode="t2i", prompt="p1", is_premium=True + ) + premium_2, _ = await queue.enqueue( + user_id=4, mode="t2i", prompt="p2", is_premium=True + ) + + order = [await queue._next_request() for _ in range(4)] # type: ignore[attr-defined] + assert order == [ + premium_1.id, + normal_1.id, + premium_2.id, + normal_2.id, + ] + + await queue.stop()