Баланс и рефералы

README.md

@@ -12,6 +12,7 @@ Telegram-бот для генерации и редактирования изо
 - **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов.
 - **Оплата**: 5 Telegram Stars за обычную очередь, 15 за премиум; админам бесплатно, возврат при ошибке.
 - **Премиум-очередь**: чередуется с обычной (одна премиум, одна обычная); если одна пуста — работает другая.
+- **Баланс и рефералы**: 3 стартовые генерации, +1 ежедневно в 00:00 МСК, +3 за первую оплату друга.
 - Глобальная очередь с одной активной задачей на пользователя.
 - Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
 - SQLite: запросы, статусы, seed, пути к файлам.
@@ -123,6 +124,18 @@ BLOCKED_PUBLIC_FIGURES=
 
 При блокировке пользователь получает «Запрос нарушает правила. Звёзды не списаны.», а запись с `user_id`, временем, промптом и причиной сохраняется в таблицу `blocked_requests`.
 
+## Баланс и рефералы
+
+- Новому пользователю начисляется 3 генерации.
+- Каждый день в 00:00 МСК — ещё 1 генерация.
+- Обычная генерация может быть оплачена 1 генерацией с баланса или 5 звёздами.
+- Премиум-очередь всегда оплачивается 15 звёздами; баланс при этом не расходуется.
+- `/balance` — текущий баланс.
+- `/invite` — реферальная ссылка.
+- За первую оплату приглашённого друга начисляется +3 генерации, реферер получает уведомление.
+
+Баланс и рефералы хранятся в таблице `users`. Пропущенные дни начисляются при следующем запуске/полуночи.
+
 ## Оплата
 
 Генерация стоит 5 Telegram Stars в обычной очереди и 15 звёзд в премиум-очереди:

bot.py

@@ -4,6 +4,8 @@ import asyncio
 import logging
 import time
 import uuid
+from datetime import datetime, timedelta, timezone
+from datetime import time as dt_time
 from pathlib import Path
 
 from telegram import (
@@ -34,12 +36,15 @@ from queue_manager import QueueEvent, QueueFull, QueueManager
 from workflows import WorkflowBuilder
 
 logger = logging.getLogger(__name__)
+MSK = timezone(timedelta(hours=3))
 
 START_TEXT = (
     "Пришли текст — сгенерирую картинку.\n"
     "Пришли фото с подписью — отредактирую его по описанию.\n"
     "Или ответь текстом на готовую картинку — возьму её как исходник.\n\n"
-    "/status — узнать состояние текущей задачи."
+    "/balance — баланс генераций.\n"
+    "/invite — пригласить друга.\n"
+    "/status — состояние текущей задачи."
 )
 BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны."
 
@@ -54,6 +59,7 @@ class MpakoBot:
         self._bot = None
         self._status_messages: dict[int, tuple[int, int]] = {}
         self._last_edit: dict[int, tuple[float, int]] = {}
+        self._daily_task: asyncio.Task[None] | None = None
 
     def run(self) -> None:
         builder = (
@@ -75,6 +81,8 @@ class MpakoBot:
         application.add_handler(CommandHandler("start", self._start))
         application.add_handler(CommandHandler("help", self._help))
         application.add_handler(CommandHandler("status", self._status))
+        application.add_handler(CommandHandler("balance", self._balance))
+        application.add_handler(CommandHandler("invite", self._invite))
         application.add_handler(
             PreCheckoutQueryHandler(self._pre_checkout)
         )
@@ -104,6 +112,10 @@ class MpakoBot:
             self._settings.comfy_ws_url,
         )
         await self._db.connect()
+        await self._db.accrue_daily_credits()
+        self._daily_task = asyncio.create_task(
+            self._daily_credits_loop(), name="daily-credits"
+        )
         enhancer = (
             PromptEnhancer(self._comfy, self._settings)
             if self._settings.prompt_enhance_enabled
@@ -121,6 +133,10 @@ class MpakoBot:
         logger.info("Бот запущен, ComfyUI: %s", self._settings.comfy_base_url)
 
     async def _post_shutdown(self, application: Application) -> None:
+        if self._daily_task is not None:
+            self._daily_task.cancel()
+            await asyncio.gather(self._daily_task, return_exceptions=True)
+            self._daily_task = None
         if self._queue is not None:
             await self._queue.stop()
         if self._comfy is not None:
@@ -128,17 +144,55 @@ class MpakoBot:
         await self._db.close()
         logger.info("Бот остановлен")
 
-    def _start_text(self) -> str:
+    async def _daily_credits_loop(self) -> None:
+        while True:
+            now = datetime.now(MSK)
+            next_midnight = datetime.combine(
+                now.date() + timedelta(days=1), dt_time.min, MSK
+            )
+            await asyncio.sleep(max(1.0, (next_midnight - now).total_seconds()))
+            try:
+                granted = await self._db.accrue_daily_credits()
+                logger.info("Ежедневные генерации начислены: %s", granted)
+            except Exception:
+                logger.exception("Ошибка ежедневного начисления генераций")
+
+    def _start_text(self, balance: int | None = None) -> str:
+        parts: list[str] = []
+        if balance is not None:
+            parts.append(
+                f"Баланс: {balance} {self._credits_word(balance)}."
+            )
         if self._settings.payments_enabled:
             normal = self._settings.stars_price
             premium = self._settings.premium_stars_price
-            return (
+            parts.append(
                 "Стоимость генерации:\n"
                 f"• обычная очередь — {normal} звёзд;\n"
-                f"• премиум-очередь — {premium} звёзд.\n\n"
+                f"• премиум-очередь — {premium} звёзд."
-                + START_TEXT
             )
-        return START_TEXT
+        parts.append(START_TEXT)
+        return "\n\n".join(parts)
+
+    @staticmethod
+    def _credits_word(count: int) -> str:
+        if count % 10 == 1 and count % 100 != 11:
+            return "генерация"
+        if count % 10 in (2, 3, 4) and count % 100 not in (12, 13, 14):
+            return "генерации"
+        return "генераций"
+
+    @staticmethod
+    def _parse_referrer(args: list[str]) -> int | None:
+        if not args:
+            return None
+        raw = args[0].strip()
+        if not raw.startswith("ref_"):
+            return None
+        try:
+            return int(raw[4:])
+        except ValueError:
+            return None
 
     async def _start(
         self, update: Update, context: ContextTypes.DEFAULT_TYPE
@@ -146,8 +200,19 @@ class MpakoBot:
         if not await self._check_access(update):
             return
         message = update.effective_message
-        if message is not None:
+        user = update.effective_user
-            await message.reply_text(self._start_text())
+        if message is None or user is None:
+            return
+        row = await self._db.ensure_user(user.id, user.username)
+        referrer_id = self._parse_referrer(context.args)
+        if (
+            referrer_id
+            and referrer_id != user.id
+            and await self._db.get_user(referrer_id) is not None
+        ):
+            await self._db.set_referrer(user.id, referrer_id)
+            row = await self._db.ensure_user(user.id, user.username)
+        await message.reply_text(self._start_text(row.credits))
 
     async def _help(
         self, update: Update, context: ContextTypes.DEFAULT_TYPE
@@ -155,8 +220,46 @@ class MpakoBot:
         if not await self._check_access(update):
             return
         message = update.effective_message
-        if message is not None:
+        user = update.effective_user
-            await message.reply_text(self._start_text())
+        if message is None or user is None:
+            return
+        row = await self._db.ensure_user(user.id, user.username)
+        await message.reply_text(self._start_text(row.credits))
+
+    async def _balance(
+        self, update: Update, context: ContextTypes.DEFAULT_TYPE
+    ) -> None:
+        if not await self._check_access(update):
+            return
+        message = update.effective_message
+        user = update.effective_user
+        if message is None or user is None:
+            return
+        row = await self._db.ensure_user(user.id, user.username)
+        await message.reply_text(
+            f"Баланс: {row.credits} {self._credits_word(row.credits)}."
+        )
+
+    async def _invite(
+        self, update: Update, context: ContextTypes.DEFAULT_TYPE
+    ) -> None:
+        if not await self._check_access(update):
+            return
+        message = update.effective_message
+        user = update.effective_user
+        if message is None or user is None or self._bot is None:
+            return
+        await self._db.ensure_user(user.id, user.username)
+        username = self._bot.username
+        if not username:
+            await message.reply_text("Не удалось получить ссылку. Попробуй позже.")
+            return
+        link = f"https://t.me/{username}?start=ref_{user.id}"
+        await message.reply_text(
+            "Пригласи друга:\n"
+            f"{link}\n\n"
+            "Когда друг впервые оплатит генерацию, ты получишь +3 генерации."
+        )
 
     async def _status(
         self, update: Update, context: ContextTypes.DEFAULT_TYPE
@@ -181,6 +284,27 @@ class MpakoBot:
             f"Задача в {label}. Перед тобой: {position} {self._people(position)}."
         )
 
+    async def _reward_referrer(self, user) -> None:
+        if not await self._db.mark_first_payment(user.id):
+            return
+        row = await self._db.get_user(user.id)
+        if row is None or row.referrer_id is None:
+            return
+        balance = await self._db.add_credits(row.referrer_id, 3)
+        if balance is None:
+            logger.warning("Не найден реферер %s", row.referrer_id)
+            return
+        text = (
+            f"Твой друг {row.display_user} сделал первую оплату. "
+            f"+3 генерации. Баланс: {balance}."
+        )
+        try:
+            await self._retry(
+                self._bot.send_message, chat_id=row.referrer_id, text=text
+            )
+        except TelegramError:
+            logger.exception("Не удалось уведомить реферера %s", row.referrer_id)
+
     async def _pre_checkout(
         self, update: Update, context: ContextTypes.DEFAULT_TYPE
     ) -> None:
@@ -216,7 +340,7 @@ class MpakoBot:
         except ValueError:
             await query.answer("Некорректный выбор.", show_alert=True)
             return
-        if tier not in {"normal", "premium"}:
+        if tier not in {"normal", "premium", "credit"}:
             await query.answer("Некорректный выбор.", show_alert=True)
             return
 
@@ -237,6 +361,39 @@ class MpakoBot:
             )
             return
 
+        if tier == "credit":
+            remaining = await self._db.consume_credit(query.from_user.id)
+            if remaining is None:
+                await query.answer(
+                    "Генерации закончились. Выбери оплату.", show_alert=True
+                )
+                return
+            await self._db.mark_pending_cancelled(payload)
+            await query.answer()
+            try:
+                await query.edit_message_text(
+                    f"Списана 1 генерация. Осталось: {remaining}."
+                )
+            except TelegramError:
+                logger.debug("Не удалось обновить сообщение выбора очереди")
+            accepted = await self._submit(
+                update,
+                mode=pending.mode,
+                prompt=pending.prompt,
+                input_path=pending.input_file_path,
+                skip_payment=True,
+                is_premium=False,
+            )
+            if not accepted:
+                await self._db.add_credits(query.from_user.id, 1)
+                try:
+                    await query.edit_message_text(
+                        "Не удалось поставить задачу. Генерация возвращена."
+                    )
+                except TelegramError:
+                    pass
+            return
+
         is_premium = tier == "premium"
         amount = (
             self._settings.premium_stars_price
@@ -309,7 +466,7 @@ class MpakoBot:
             await message.reply_text("Этот платёж уже обработан.")
             return
 
-        await self._submit(
+        accepted = await self._submit(
             update,
             mode=pending.mode,
             prompt=pending.prompt,
@@ -319,6 +476,8 @@ class MpakoBot:
             payment_payload=payload,
             is_premium=pending.is_premium,
         )
+        if accepted:
+            await self._reward_referrer(user)
 
     async def _text(
         self, update: Update, context: ContextTypes.DEFAULT_TYPE
@@ -471,6 +630,8 @@ class MpakoBot:
             return False
 
         is_admin = user.id in self._settings.admin_user_ids
+        user_row = await self._db.ensure_user(user.id, user.username)
+        credits = user_row.credits
         has_person = False
         if (
             mode == "i2i"
@@ -527,8 +688,16 @@ class MpakoBot:
                 [
                     [
                         InlineKeyboardButton(
-                            f"Обычная — {self._settings.stars_price} звёзд",
+                            (
-                            callback_data=f"queue:{payload}:normal",
+                                f"Обычная — 1 генерация (осталось {credits})"
+                                if credits > 0
+                                else f"Обычная — {self._settings.stars_price} звёзд"
+                            ),
+                            callback_data=(
+                                f"queue:{payload}:credit"
+                                if credits > 0
+                                else f"queue:{payload}:normal"
+                            ),
                         )
                     ],
                     [
@@ -540,6 +709,7 @@ class MpakoBot:
                 ]
             )
             text = (
+                f"Баланс: {credits} {self._credits_word(credits)}.\n"
                 f"Обычная очередь: {counts['normal']} чел.\n"
                 f"Премиум-очередь: {counts['premium']} чел.\n\n"
                 "Выбери очередь:"

checklist.md

@@ -6,10 +6,10 @@
 
 ## Текущий статус
 
-- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15) и оплата Stars (16).
+- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15), оплата Stars (16) и баланс/рефералы (17).
 - Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история.
-- 49 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, вспомогательные функции ComfyUI.
+- 50 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, вспомогательные функции ComfyUI.
-- Не проверено: реальная оплата Stars в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото.
+- Не проверено: реальная оплата Stars и реферальная выплата в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото.
 
 ## 0. Допущения и решения по умолчанию
 
@@ -341,6 +341,16 @@ ON requests(user_id) WHERE status IN ('queued', 'running');
 
 **DoD:** не-админ получает счёт на 5 звёзд, после оплаты задача ставится в очередь; при сбое звёзды возвращаются.
 
+## 17. Баланс и рефералы
+
+- [x] `users`: 3 стартовые генерации, ежедневное +1 в 00:00 МСК.
+- [x] Списание 1 генерации за обычную очередь, премиум — только звёзды.
+- [x] `/balance`, `/invite`, реферальная ссылка `?start=ref_<id>`.
+- [x] +3 генерации рефереру за первую оплату друга и уведомление.
+- [x] Тесты баланса, ежедневного начисления и рефералов.
+
+**DoD:** новый пользователь получает 3 генерации, ежедневно +1, друг после первой оплаты приносит рефереру +3.
+
 ## Порядок работ
 
 1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow.

db.py

@@ -4,7 +4,7 @@ import asyncio
 import logging
 import sqlite3
 from dataclasses import dataclass
-from datetime import UTC, datetime
+from datetime import UTC, date, datetime, timedelta, timezone
 from pathlib import Path
 from typing import Any
 
@@ -13,12 +13,17 @@ import aiosqlite
 logger = logging.getLogger(__name__)
 
 ACTIVE_STATUSES = ("queued", "running")
+MSK = timezone(timedelta(hours=3))
 
 
 def utc_now() -> str:
     return datetime.now(UTC).isoformat()
 
 
+def msk_today() -> date:
+    return datetime.now(MSK).date()
+
+
 class UserAlreadyQueued(RuntimeError):
     def __init__(self, user_id: int) -> None:
         super().__init__(f"У пользователя {user_id} уже есть активный запрос")
@@ -107,6 +112,33 @@ class BlockedRequest:
         )
 
 
+@dataclass(frozen=True)
+class User:
+    user_id: int
+    username: str | None
+    credits: int
+    last_daily_credit: str
+    referrer_id: int | None
+    first_paid_at: str | None
+    created_at: str
+
+    @property
+    def display_user(self) -> str:
+        return f"@{self.username}" if self.username else str(self.user_id)
+
+    @classmethod
+    def from_row(cls, row: aiosqlite.Row) -> User:
+        return cls(
+            user_id=int(row["user_id"]),
+            username=row["username"],
+            credits=int(row["credits"]),
+            last_daily_credit=row["last_daily_credit"],
+            referrer_id=row["referrer_id"],
+            first_paid_at=row["first_paid_at"],
+            created_at=row["created_at"],
+        )
+
+
 @dataclass(frozen=True)
 class PendingPayment:
     payload: str
@@ -202,6 +234,19 @@ CREATE TABLE IF NOT EXISTS pending_payments (
 
 CREATE INDEX IF NOT EXISTS idx_pending_user_time
     ON pending_payments(user_id, timestamp);
+
+CREATE TABLE IF NOT EXISTS users (
+    user_id            INTEGER PRIMARY KEY,
+    username           TEXT,
+    credits            INTEGER NOT NULL DEFAULT 3,
+    last_daily_credit  TEXT    NOT NULL,
+    referrer_id        INTEGER,
+    first_paid_at      TEXT,
+    created_at         TEXT    NOT NULL
+);
+
+CREATE INDEX IF NOT EXISTS idx_users_referrer
+    ON users(referrer_id);
 """
 
 
@@ -250,6 +295,16 @@ class Database:
             await conn.execute(
                 "ALTER TABLE pending_payments ADD COLUMN is_premium INTEGER NOT NULL DEFAULT 0"
             )
+        await conn.execute(
+            """
+            INSERT OR IGNORE INTO users
+                (user_id, username, credits, last_daily_credit,
+                 referrer_id, first_paid_at, created_at)
+            SELECT user_id, NULL, 3, ?, NULL, NULL, ?
+            FROM (SELECT DISTINCT user_id FROM requests)
+            """,
+            (msk_today().isoformat(), utc_now()),
+        )
 
     async def close(self) -> None:
         if self._conn is not None:
@@ -327,6 +382,115 @@ class Database:
             row = await cursor.fetchone()
         return Request.from_row(row) if row is not None else None
 
+    async def get_user(self, user_id: int) -> User | None:
+        conn = self._connection()
+        async with conn.execute(
+            "SELECT * FROM users WHERE user_id = ?", (user_id,)
+        ) as cursor:
+            row = await cursor.fetchone()
+        return User.from_row(row) if row is not None else None
+
+    async def ensure_user(
+        self, user_id: int, username: str | None = None
+    ) -> User:
+        conn = self._connection()
+        today = msk_today().isoformat()
+        async with self._write_lock:
+            await conn.execute(
+                """
+                INSERT OR IGNORE INTO users
+                    (user_id, username, credits, last_daily_credit, created_at)
+                VALUES (?, ?, 3, ?, ?)
+                """,
+                (user_id, username, today, utc_now()),
+            )
+            if username:
+                await conn.execute(
+                    "UPDATE users SET username = ? WHERE user_id = ?",
+                    (username, user_id),
+                )
+            await conn.commit()
+        user = await self.get_user(user_id)
+        if user is None:
+            raise RuntimeError("Не удалось прочитать пользователя")
+        return user
+
+    async def add_credits(self, user_id: int, amount: int) -> int | None:
+        conn = self._connection()
+        async with self._write_lock:
+            cursor = await conn.execute(
+                "UPDATE users SET credits = credits + ? WHERE user_id = ?",
+                (amount, user_id),
+            )
+            await conn.commit()
+            if cursor.rowcount == 0:
+                return None
+        user = await self.get_user(user_id)
+        return user.credits if user is not None else None
+
+    async def consume_credit(self, user_id: int) -> int | None:
+        conn = self._connection()
+        async with self._write_lock:
+            cursor = await conn.execute(
+                "UPDATE users SET credits = credits - 1"
+                " WHERE user_id = ? AND credits > 0",
+                (user_id,),
+            )
+            await conn.commit()
+            if cursor.rowcount == 0:
+                return None
+        user = await self.get_user(user_id)
+        return user.credits if user is not None else None
+
+    async def set_referrer(self, user_id: int, referrer_id: int) -> bool:
+        conn = self._connection()
+        async with self._write_lock:
+            cursor = await conn.execute(
+                """
+                UPDATE users SET referrer_id = ?
+                WHERE user_id = ? AND referrer_id IS NULL AND user_id != ?
+                """,
+                (referrer_id, user_id, referrer_id),
+            )
+            await conn.commit()
+            return cursor.rowcount > 0
+
+    async def mark_first_payment(self, user_id: int) -> bool:
+        conn = self._connection()
+        async with self._write_lock:
+            cursor = await conn.execute(
+                "UPDATE users SET first_paid_at = ?"
+                " WHERE user_id = ? AND first_paid_at IS NULL",
+                (utc_now(), user_id),
+            )
+            await conn.commit()
+            return cursor.rowcount > 0
+
+    async def accrue_daily_credits(self, today: date | None = None) -> int:
+        conn = self._connection()
+        current = today or msk_today()
+        current_text = current.isoformat()
+        async with self._write_lock:
+            async with conn.execute(
+                "SELECT user_id, last_daily_credit FROM users"
+                " WHERE last_daily_credit < ?",
+                (current_text,),
+            ) as cursor:
+                rows = await cursor.fetchall()
+            for row in rows:
+                try:
+                    last = date.fromisoformat(row["last_daily_credit"])
+                    days = max(1, (current - last).days)
+                except ValueError:
+                    days = 1
+                await conn.execute(
+                    "UPDATE users SET credits = credits + ?,"
+                    " last_daily_credit = ? WHERE user_id = ?",
+                    (days, current_text, int(row["user_id"])),
+                )
+            await conn.commit()
+        return len(rows)
+
     async def list_queued(self) -> list[Request]:
         conn = self._connection()
         async with conn.execute(

tests/test_db.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import sqlite3
+from datetime import date, timedelta
 
 import pytest
 import pytest_asyncio
@@ -242,6 +243,35 @@ async def test_pending_payment_lifecycle(db: Database) -> None:
     assert loaded_request.payment_refunded is True
 
 
+async def test_user_credits_referral_and_daily_accrual(db: Database) -> None:
+    user = await db.ensure_user(1, "alice")
+    assert user.credits == 3
+    assert user.last_daily_credit
+
+    assert await db.add_credits(1, 2) == 5
+    assert await db.consume_credit(1) == 4
+    assert await db.consume_credit(1) == 3
+
+    friend = await db.ensure_user(2, "bob")
+    assert friend.credits == 3
+    assert await db.set_referrer(2, 1) is True
+    assert await db.set_referrer(2, 1) is False
+    assert await db.mark_first_payment(2) is True
+    assert await db.mark_first_payment(2) is False
+
+    friend = await db.get_user(2)
+    assert friend is not None
+    assert friend.referrer_id == 1
+    assert friend.first_paid_at is not None
+
+    future = date.today() + timedelta(days=2)
+    assert await db.accrue_daily_credits(future) >= 2
+    user = await db.get_user(1)
+    assert user is not None
+    assert user.credits == 5
+    assert await db.accrue_daily_credits(future) == 0
+
+
 async def test_migrates_legacy_schema_without_username(tmp_path) -> None:
     path = tmp_path / "legacy.db"
     connection = sqlite3.connect(path)