Оплата генераций звёздами

Счёт в XTR, возврат при ошибке, бесплатный доступ для админов.
This commit is contained in:
MpakobecEx 2026-10-13 21:38:00 +05:00
parent e90e891c3b
commit 960672818f
8 changed files with 445 additions and 14 deletions

View File

@ -40,6 +40,10 @@ MODERATION_ENABLED=true
# Дополнительные имена публичных людей через запятую # Дополнительные имена публичных людей через запятую
BLOCKED_PUBLIC_FIGURES= BLOCKED_PUBLIC_FIGURES=
# Оплата Telegram Stars
PAYMENTS_ENABLED=true
STARS_PRICE=5
# Веб-история: только локальная сеть, без авторизации # Веб-история: только локальная сеть, без авторизации
WEB_BIND_IP=192.168.31.158 WEB_BIND_IP=192.168.31.158
WEB_PAGE_SIZE=50 WEB_PAGE_SIZE=50

View File

@ -10,6 +10,7 @@ Telegram-бот для генерации и редактирования изо
- **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений. - **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений.
- **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI. - **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI.
- **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов. - **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов.
- **Оплата**: 5 Telegram Stars за генерацию, админам бесплатно, возврат при ошибке.
- Глобальная очередь с одной активной задачей на пользователя. - Глобальная очередь с одной активной задачей на пользователя.
- Позиция в очереди и прогресс генерации через WebSocket ComfyUI. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
- SQLite: запросы, статусы, seed, пути к файлам. - SQLite: запросы, статусы, seed, пути к файлам.
@ -121,6 +122,22 @@ BLOCKED_PUBLIC_FIGURES=
При блокировке пользователь получает «Запрос нарушает правила. Звёзды не списаны.», а запись с `user_id`, временем, промптом и причиной сохраняется в таблицу `blocked_requests`. При блокировке пользователь получает «Запрос нарушает правила. Звёзды не списаны.», а запись с `user_id`, временем, промптом и причиной сохраняется в таблицу `blocked_requests`.
## Оплата
Генерация стоит 5 Telegram Stars:
```dotenv
PAYMENTS_ENABLED=true
STARS_PRICE=5
```
- Админы из `ADMIN_USER_IDS` генерируют бесплатно.
- Не-админу выставляется счёт через Telegram Stars (`XTR`).
- Фильтрация промпта выполняется до выставления счёта.
- Если поставить задачу в очередь не удалось, звёзды возвращаются.
- При ошибке генерации звёзды также возвращаются.
- Оплата и `telegram_payment_charge_id` сохраняются в `pending_payments` и `requests`.
## Веб-история ## Веб-история
После `docker compose up -d` доступна страница: После `docker compose up -d` доступна страница:

217
bot.py
View File

@ -6,7 +6,7 @@ import time
import uuid import uuid
from pathlib import Path from pathlib import Path
from telegram import InputFile, Update from telegram import InputFile, LabeledPrice, Update
from telegram.constants import ChatAction, ChatType from telegram.constants import ChatAction, ChatType
from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut
from telegram.ext import ( from telegram.ext import (
@ -14,13 +14,14 @@ from telegram.ext import (
CommandHandler, CommandHandler,
ContextTypes, ContextTypes,
MessageHandler, MessageHandler,
PreCheckoutQueryHandler,
filters, filters,
) )
from comfy_client import ComfyClient from comfy_client import ComfyClient
from config import Settings, setup_logging from config import Settings, setup_logging
from db import Database, Request, UserAlreadyQueued from db import Database, Request, UserAlreadyQueued
from moderation import PromptBlocked, detect_person from moderation import PromptBlocked, detect_person, is_prompt_allowed
from prompt_enhancer import PromptEnhancer from prompt_enhancer import PromptEnhancer
from queue_manager import QueueEvent, QueueFull, QueueManager from queue_manager import QueueEvent, QueueFull, QueueManager
from workflows import WorkflowBuilder from workflows import WorkflowBuilder
@ -67,6 +68,12 @@ class MpakoBot:
application.add_handler(CommandHandler("start", self._start)) application.add_handler(CommandHandler("start", self._start))
application.add_handler(CommandHandler("help", self._help)) application.add_handler(CommandHandler("help", self._help))
application.add_handler(CommandHandler("status", self._status)) application.add_handler(CommandHandler("status", self._status))
application.add_handler(
PreCheckoutQueryHandler(self._pre_checkout)
)
application.add_handler(
MessageHandler(filters.SUCCESSFUL_PAYMENT, self._successful_payment)
)
application.add_handler(MessageHandler(filters.PHOTO, self._photo)) application.add_handler(MessageHandler(filters.PHOTO, self._photo))
application.add_handler( application.add_handler(
MessageHandler(filters.TEXT & ~filters.COMMAND, self._text) MessageHandler(filters.TEXT & ~filters.COMMAND, self._text)
@ -111,6 +118,15 @@ class MpakoBot:
await self._db.close() await self._db.close()
logger.info("Бот остановлен") logger.info("Бот остановлен")
def _start_text(self) -> str:
text = START_TEXT
if self._settings.payments_enabled:
text += (
f"\n\nСтоимость генерации: {self._settings.stars_price} звёзд. "
"Админам бесплатно."
)
return text
async def _start( async def _start(
self, update: Update, context: ContextTypes.DEFAULT_TYPE self, update: Update, context: ContextTypes.DEFAULT_TYPE
) -> None: ) -> None:
@ -118,7 +134,7 @@ class MpakoBot:
return return
message = update.effective_message message = update.effective_message
if message is not None: if message is not None:
await message.reply_text(START_TEXT) await message.reply_text(self._start_text())
async def _help( async def _help(
self, update: Update, context: ContextTypes.DEFAULT_TYPE self, update: Update, context: ContextTypes.DEFAULT_TYPE
@ -127,7 +143,7 @@ class MpakoBot:
return return
message = update.effective_message message = update.effective_message
if message is not None: if message is not None:
await message.reply_text(START_TEXT) await message.reply_text(self._start_text())
async def _status( async def _status(
self, update: Update, context: ContextTypes.DEFAULT_TYPE self, update: Update, context: ContextTypes.DEFAULT_TYPE
@ -150,6 +166,76 @@ class MpakoBot:
f"Задача в очереди. Перед тобой: {position} {self._people(position)}." f"Задача в очереди. Перед тобой: {position} {self._people(position)}."
) )
async def _pre_checkout(
self, update: Update, context: ContextTypes.DEFAULT_TYPE
) -> None:
query = update.pre_checkout_query
if query is None:
return
pending = await self._db.get_pending_payment(query.invoice_payload)
allowed = (
pending is not None
and pending.user_id == query.from_user.id
and pending.status == "pending"
and query.currency == "XTR"
and query.total_amount == pending.amount
)
try:
await query.answer(
ok=allowed,
error_message=None
if allowed
else "Счёт устарел. Отправь запрос заново.",
)
except TelegramError:
logger.exception("Не удалось ответить на pre_checkout_query")
async def _successful_payment(
self, update: Update, context: ContextTypes.DEFAULT_TYPE
) -> None:
message = update.effective_message
user = update.effective_user
payment = message.successful_payment if message is not None else None
if message is None or user is None or payment is None:
return
payload = payment.invoice_payload
pending = await self._db.get_pending_payment(payload)
if pending is None or pending.user_id != user.id:
logger.error(
"Платёж %s не найден или принадлежит другому пользователю",
payload,
)
return
if pending.status == "paid":
await message.reply_text("Этот платёж уже обработан.")
return
if pending.status != "pending":
return
charge_id = payment.telegram_payment_charge_id
if payment.currency != "XTR" or payment.total_amount != pending.amount:
refunded = await self._refund_stars(user.id, charge_id)
if refunded:
await self._db.mark_pending_refunded(payload)
note = " Звёзды возвращены." if refunded else ""
await message.reply_text("Платёж не принят." + note)
return
if not await self._db.mark_pending_paid(payload, charge_id):
await message.reply_text("Этот платёж уже обработан.")
return
await self._submit(
update,
mode=pending.mode,
prompt=pending.prompt,
input_path=pending.input_file_path,
skip_payment=True,
payment_charge_id=charge_id,
payment_payload=payload,
)
async def _text( async def _text(
self, update: Update, context: ContextTypes.DEFAULT_TYPE self, update: Update, context: ContextTypes.DEFAULT_TYPE
) -> None: ) -> None:
@ -290,6 +376,9 @@ class MpakoBot:
mode: str, mode: str,
prompt: str, prompt: str,
input_path: str | None, input_path: str | None,
skip_payment: bool = False,
payment_charge_id: str | None = None,
payment_payload: str | None = None,
) -> bool: ) -> bool:
message = update.effective_message message = update.effective_message
user = update.effective_user user = update.effective_user
@ -306,6 +395,59 @@ class MpakoBot:
): ):
has_person = await asyncio.to_thread(detect_person, input_path) has_person = await asyncio.to_thread(detect_person, input_path)
if (
not skip_payment
and not is_admin
and self._settings.payments_enabled
):
allowed, reason = is_prompt_allowed(
prompt,
mode=mode,
has_person=has_person,
extra_public_figures=self._settings.blocked_public_figures,
)
if not allowed:
await self._db.log_blocked(
user_id=user.id,
prompt=prompt,
reason=reason or "unknown",
mode=mode,
has_person=has_person,
)
await message.reply_text(BLOCKED_TEXT)
return False
payload = uuid.uuid4().hex
amount = self._settings.stars_price
await self._db.create_pending_payment(
payload=payload,
user_id=user.id,
mode=mode,
prompt=prompt,
username=user.username,
input_file_path=input_path,
amount=amount,
)
try:
await self._retry(
self._bot.send_invoice,
chat_id=user.id,
title="Генерация изображения",
description="Одна генерация или редактирование изображения.",
payload=payload,
provider_token="",
currency="XTR",
prices=[LabeledPrice(label="Генерация", amount=amount)],
)
except TelegramError:
logger.exception("Не удалось выставить счёт пользователю %s", user.id)
await self._db.mark_pending_cancelled(payload)
await message.reply_text(
"Не удалось выставить счёт. Попробуй позже."
)
return False
return True
try: try:
request, position = await self._queue.enqueue( request, position = await self._queue.enqueue(
user_id=user.id, user_id=user.id,
@ -315,11 +457,18 @@ class MpakoBot:
input_file_path=input_path, input_file_path=input_path,
has_person=has_person, has_person=has_person,
is_admin=is_admin, is_admin=is_admin,
payment_charge_id=payment_charge_id,
) )
except PromptBlocked: except PromptBlocked:
await message.reply_text(BLOCKED_TEXT) note = await self._refund_if_paid(
user.id, payment_charge_id, payment_payload
)
await message.reply_text(BLOCKED_TEXT + note)
return False return False
except UserAlreadyQueued: except UserAlreadyQueued:
note = await self._refund_if_paid(
user.id, payment_charge_id, payment_payload
)
active = await self._db.get_active_by_user(user.id) active = await self._db.get_active_by_user(user.id)
text = "Ты уже в очереди." text = "Ты уже в очереди."
if active is not None: if active is not None:
@ -328,16 +477,27 @@ class MpakoBot:
else: else:
ahead = await self._queue.position(active.id) ahead = await self._queue.position(active.id)
text += f" Перед тобой: {ahead} {self._people(ahead)}." text += f" Перед тобой: {ahead} {self._people(ahead)}."
await message.reply_text(text) await message.reply_text(text + note)
return False return False
except QueueFull as exc: except QueueFull as exc:
await message.reply_text(exc.user_message) note = await self._refund_if_paid(
user.id, payment_charge_id, payment_payload
)
await message.reply_text(exc.user_message + note)
return False return False
except Exception: except Exception:
logger.exception("Не удалось поставить задачу пользователю %s", user.id) logger.exception("Не удалось поставить задачу пользователю %s", user.id)
await message.reply_text("Не удалось поставить задачу. Попробуй позже.") note = await self._refund_if_paid(
user.id, payment_charge_id, payment_payload
)
await message.reply_text(
"Не удалось поставить задачу. Попробуй позже." + note
)
return False return False
if payment_payload:
await self._db.attach_pending_request(payment_payload, request.id)
try: try:
status_message = await self._retry( status_message = await self._retry(
message.reply_text, self._ack_text(position) message.reply_text, self._ack_text(position)
@ -367,6 +527,35 @@ class MpakoBot:
) )
return True return True
async def _refund_if_paid(
self,
user_id: int,
charge_id: str | None,
payload: str | None,
) -> str:
if not charge_id:
return ""
refunded = await self._refund_stars(user_id, charge_id)
if refunded and payload:
await self._db.mark_pending_refunded(payload)
if refunded:
return " Звёзды возвращены."
return " Не удалось вернуть звёзды, напиши админу."
async def _refund_stars(self, user_id: int, charge_id: str) -> bool:
if self._bot is None:
return False
try:
return await self._retry(
self._bot.refund_star_payment,
user_id=user_id,
telegram_payment_charge_id=charge_id,
attempts=2,
)
except TelegramError:
logger.exception("Не удалось вернуть звёзды пользователю %s", user_id)
return False
def _ack_text(self, position: int) -> str: def _ack_text(self, position: int) -> str:
if position <= 0: if position <= 0:
return "Принято. Ты первый, начинаю…" return "Принято. Ты первый, начинаю…"
@ -513,7 +702,17 @@ class MpakoBot:
async def _finish_failed(self, request: Request, error: str) -> None: async def _finish_failed(self, request: Request, error: str) -> None:
target = self._status_messages.get(request.id) target = self._status_messages.get(request.id)
chat_id = target[0] if target is not None else request.user_id chat_id = target[0] if target is not None else request.user_id
text = f"Не получилось: {error}" refund_note = ""
if request.payment_charge_id and not request.payment_refunded:
refunded = await self._refund_stars(
request.user_id, request.payment_charge_id
)
if refunded:
await self._db.set_refunded(request.id)
refund_note = " Звёзды возвращены."
else:
refund_note = " Не удалось вернуть звёзды, напиши админу."
text = f"Не получилось: {error}{refund_note}"
if target is not None and self._bot is not None: if target is not None and self._bot is not None:
try: try:
await self._retry( await self._retry(

View File

@ -6,10 +6,10 @@
## Текущий статус ## Текущий статус
- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14) и модерация (15). - Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15) и оплата Stars (16).
- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история. - Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история.
- 45 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, вспомогательные функции ComfyUI. - 46 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, вспомогательные функции ComfyUI.
- Не проверено: очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. - Не проверено: реальная оплата Stars в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото.
## 0. Допущения и решения по умолчанию ## 0. Допущения и решения по умолчанию
@ -329,6 +329,17 @@ ON requests(user_id) WHERE status IN ('queued', 'running');
**DoD:** небезопасный промпт не попадает в очередь, причина пишется в БД, админ может отправлять без фильтра. **DoD:** небезопасный промпт не попадает в очередь, причина пишется в БД, админ может отправлять без фильтра.
## 16. Оплата Stars
- [x] `PAYMENTS_ENABLED` / `STARS_PRICE`, админы бесплатно.
- [x] Invoice `XTR` + `pre_checkout_query` + `successful_payment`.
- [x] `pending_payments`: payload, сумма, статус, charge_id, request_id.
- [x] Возврат звёзд при ошибке постановки в очередь и генерации.
- [x] `payment_charge_id` и `payment_refunded` в `requests`.
- [x] Тесты жизненного цикла платежа.
**DoD:** не-админ получает счёт на 5 звёзд, после оплаты задача ставится в очередь; при сбое звёзды возвращаются.
## Порядок работ ## Порядок работ
1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow. 1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow.

View File

@ -94,6 +94,8 @@ class Settings:
max_photo_bytes: int = 15 * 1024 * 1024 max_photo_bytes: int = 15 * 1024 * 1024
moderation_enabled: bool = True moderation_enabled: bool = True
blocked_public_figures: frozenset[str] = field(default_factory=frozenset) blocked_public_figures: frozenset[str] = field(default_factory=frozenset)
payments_enabled: bool = True
stars_price: int = 5
web_page_size: int = 50 web_page_size: int = 50
web_bind_ip: str = "127.0.0.1" web_bind_ip: str = "127.0.0.1"
@ -152,6 +154,8 @@ class Settings:
max_photo_bytes=max(1, _int("MAX_PHOTO_BYTES", 15 * 1024 * 1024)), max_photo_bytes=max(1, _int("MAX_PHOTO_BYTES", 15 * 1024 * 1024)),
moderation_enabled=_bool("MODERATION_ENABLED", True), moderation_enabled=_bool("MODERATION_ENABLED", True),
blocked_public_figures=_csv("BLOCKED_PUBLIC_FIGURES"), blocked_public_figures=_csv("BLOCKED_PUBLIC_FIGURES"),
payments_enabled=_bool("PAYMENTS_ENABLED", True),
stars_price=max(1, _int("STARS_PRICE", 5)),
web_page_size=max(10, min(200, _int("WEB_PAGE_SIZE", 50))), web_page_size=max(10, min(200, _int("WEB_PAGE_SIZE", 50))),
web_bind_ip=os.getenv("WEB_BIND_IP", "127.0.0.1").strip() web_bind_ip=os.getenv("WEB_BIND_IP", "127.0.0.1").strip()
or "127.0.0.1", or "127.0.0.1",

153
db.py
View File

@ -37,6 +37,8 @@ class Request:
output_file_path: str | None output_file_path: str | None
status: str status: str
comfy_prompt_id: str | None comfy_prompt_id: str | None
payment_charge_id: str | None
payment_refunded: bool
mode: str mode: str
seed: int | None seed: int | None
error_text: str | None error_text: str | None
@ -59,6 +61,8 @@ class Request:
output_file_path=row["output_file_path"], output_file_path=row["output_file_path"],
status=row["status"], status=row["status"],
comfy_prompt_id=row["comfy_prompt_id"], comfy_prompt_id=row["comfy_prompt_id"],
payment_charge_id=row["payment_charge_id"],
payment_refunded=bool(row["payment_refunded"]),
mode=row["mode"], mode=row["mode"],
seed=row["seed"], seed=row["seed"],
error_text=row["error_text"], error_text=row["error_text"],
@ -101,6 +105,37 @@ class BlockedRequest:
) )
@dataclass(frozen=True)
class PendingPayment:
payload: str
user_id: int
timestamp: str
mode: str
prompt: str
username: str | None
input_file_path: str | None
amount: int
status: str
telegram_payment_charge_id: str | None
request_id: int | None
@classmethod
def from_row(cls, row: aiosqlite.Row) -> PendingPayment:
return cls(
payload=row["payload"],
user_id=row["user_id"],
timestamp=row["timestamp"],
mode=row["mode"],
prompt=row["prompt"],
username=row["username"],
input_file_path=row["input_file_path"],
amount=int(row["amount"]),
status=row["status"],
telegram_payment_charge_id=row["telegram_payment_charge_id"],
request_id=row["request_id"],
)
SCHEMA = """ SCHEMA = """
CREATE TABLE IF NOT EXISTS requests ( CREATE TABLE IF NOT EXISTS requests (
id INTEGER PRIMARY KEY AUTOINCREMENT, id INTEGER PRIMARY KEY AUTOINCREMENT,
@ -114,6 +149,8 @@ CREATE TABLE IF NOT EXISTS requests (
status TEXT NOT NULL DEFAULT 'queued' status TEXT NOT NULL DEFAULT 'queued'
CHECK (status IN ('queued', 'running', 'done', 'failed')), CHECK (status IN ('queued', 'running', 'done', 'failed')),
comfy_prompt_id TEXT, comfy_prompt_id TEXT,
payment_charge_id TEXT,
payment_refunded INTEGER NOT NULL DEFAULT 0,
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')), mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
seed INTEGER, seed INTEGER,
error_text TEXT, error_text TEXT,
@ -141,6 +178,24 @@ CREATE TABLE IF NOT EXISTS blocked_requests (
CREATE INDEX IF NOT EXISTS idx_blocked_user_time CREATE INDEX IF NOT EXISTS idx_blocked_user_time
ON blocked_requests(user_id, id); ON blocked_requests(user_id, id);
CREATE TABLE IF NOT EXISTS pending_payments (
payload TEXT PRIMARY KEY,
user_id INTEGER NOT NULL,
timestamp TEXT NOT NULL,
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
prompt TEXT NOT NULL,
username TEXT,
input_file_path TEXT,
amount INTEGER NOT NULL,
status TEXT NOT NULL DEFAULT 'pending'
CHECK (status IN ('pending', 'paid', 'cancelled', 'refunded')),
telegram_payment_charge_id TEXT,
request_id INTEGER
);
CREATE INDEX IF NOT EXISTS idx_pending_user_time
ON pending_payments(user_id, timestamp);
""" """
@ -173,6 +228,12 @@ class Database:
await conn.execute("ALTER TABLE requests ADD COLUMN username TEXT") await conn.execute("ALTER TABLE requests ADD COLUMN username TEXT")
if "enhanced_prompt" not in columns: if "enhanced_prompt" not in columns:
await conn.execute("ALTER TABLE requests ADD COLUMN enhanced_prompt TEXT") await conn.execute("ALTER TABLE requests ADD COLUMN enhanced_prompt TEXT")
if "payment_charge_id" not in columns:
await conn.execute("ALTER TABLE requests ADD COLUMN payment_charge_id TEXT")
if "payment_refunded" not in columns:
await conn.execute(
"ALTER TABLE requests ADD COLUMN payment_refunded INTEGER NOT NULL DEFAULT 0"
)
async def close(self) -> None: async def close(self) -> None:
if self._conn is not None: if self._conn is not None:
@ -193,6 +254,7 @@ class Database:
prompt_text: str, prompt_text: str,
input_file_path: str | None, input_file_path: str | None,
seed: int, seed: int,
payment_charge_id: str | None = None,
) -> Request: ) -> Request:
conn = self._connection() conn = self._connection()
async with self._write_lock: async with self._write_lock:
@ -201,8 +263,8 @@ class Database:
""" """
INSERT INTO requests INSERT INTO requests
(user_id, username, timestamp, prompt_text, (user_id, username, timestamp, prompt_text,
input_file_path, status, mode, seed) input_file_path, status, mode, seed, payment_charge_id)
VALUES (?, ?, ?, ?, ?, 'queued', ?, ?) VALUES (?, ?, ?, ?, ?, 'queued', ?, ?, ?)
""", """,
( (
user_id, user_id,
@ -212,6 +274,7 @@ class Database:
input_file_path, input_file_path,
mode, mode,
seed, seed,
payment_charge_id,
), ),
) )
request_id = cursor.lastrowid request_id = cursor.lastrowid
@ -389,6 +452,92 @@ class Database:
(prompt, request_id), (prompt, request_id),
) )
async def set_refunded(self, request_id: int) -> None:
await self._update(
"UPDATE requests SET payment_refunded = 1 WHERE id = ?",
(request_id,),
)
async def create_pending_payment(
self,
*,
payload: str,
user_id: int,
mode: str,
prompt: str,
username: str | None,
input_file_path: str | None,
amount: int,
) -> PendingPayment:
conn = self._connection()
async with self._write_lock:
await conn.execute(
"""
INSERT INTO pending_payments
(payload, user_id, timestamp, mode, prompt, username,
input_file_path, amount, status)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending')
""",
(
payload,
user_id,
utc_now(),
mode,
prompt,
username,
input_file_path,
amount,
),
)
await conn.commit()
pending = await self.get_pending_payment(payload)
if pending is None:
raise RuntimeError("Не удалось прочитать счёт")
return pending
async def get_pending_payment(self, payload: str) -> PendingPayment | None:
conn = self._connection()
async with conn.execute(
"SELECT * FROM pending_payments WHERE payload = ?", (payload,)
) as cursor:
row = await cursor.fetchone()
return PendingPayment.from_row(row) if row is not None else None
async def mark_pending_paid(self, payload: str, charge_id: str) -> bool:
conn = self._connection()
async with self._write_lock:
cursor = await conn.execute(
"""
UPDATE pending_payments
SET status = 'paid', telegram_payment_charge_id = ?
WHERE payload = ? AND status = 'pending'
""",
(charge_id, payload),
)
await conn.commit()
return cursor.rowcount > 0
async def mark_pending_refunded(self, payload: str) -> None:
await self._update(
"UPDATE pending_payments SET status = 'refunded' WHERE payload = ?",
(payload,),
)
async def mark_pending_cancelled(self, payload: str) -> None:
await self._update(
"""
UPDATE pending_payments SET status = 'cancelled'
WHERE payload = ? AND status = 'pending'
""",
(payload,),
)
async def attach_pending_request(self, payload: str, request_id: int) -> None:
await self._update(
"UPDATE pending_payments SET request_id = ? WHERE payload = ?",
(request_id, payload),
)
async def set_done(self, request_id: int, output_file_path: str) -> None: async def set_done(self, request_id: int, output_file_path: str) -> None:
await self._update( await self._update(
""" """

View File

@ -107,6 +107,7 @@ class QueueManager:
username: str | None = None, username: str | None = None,
has_person: bool = False, has_person: bool = False,
is_admin: bool = False, is_admin: bool = False,
payment_charge_id: str | None = None,
) -> tuple[Request, int]: ) -> tuple[Request, int]:
async with self._lock: async with self._lock:
if self._stopping: if self._stopping:
@ -141,6 +142,7 @@ class QueueManager:
prompt_text=prompt, prompt_text=prompt,
input_file_path=input_file_path, input_file_path=input_file_path,
seed=seed, seed=seed,
payment_charge_id=payment_charge_id,
) )
self._queue.append(request.id) self._queue.append(request.id)
self._active_users.add(user_id) self._active_users.add(user_id)

View File

@ -165,6 +165,51 @@ async def test_history_filters_and_stats(db: Database) -> None:
assert alice.display_user == "@alice" assert alice.display_user == "@alice"
async def test_pending_payment_lifecycle(db: Database) -> None:
pending = await db.create_pending_payment(
payload="payload-1",
user_id=1,
mode="t2i",
prompt="apple",
username="alice",
input_file_path=None,
amount=5,
)
assert pending.status == "pending"
assert pending.amount == 5
assert await db.mark_pending_paid("payload-1", "charge-1") is True
loaded = await db.get_pending_payment("payload-1")
assert loaded is not None
assert loaded.status == "paid"
assert loaded.telegram_payment_charge_id == "charge-1"
await db.attach_pending_request("payload-1", 42)
loaded = await db.get_pending_payment("payload-1")
assert loaded is not None
assert loaded.request_id == 42
await db.mark_pending_refunded("payload-1")
loaded = await db.get_pending_payment("payload-1")
assert loaded is not None
assert loaded.status == "refunded"
request = await db.create_request(
user_id=2,
mode="t2i",
prompt_text="paid apple",
input_file_path=None,
seed=1,
payment_charge_id="charge-2",
)
assert request.payment_charge_id == "charge-2"
assert request.payment_refunded is False
await db.set_refunded(request.id)
loaded_request = await db.get_request(request.id)
assert loaded_request is not None
assert loaded_request.payment_refunded is True
async def test_migrates_legacy_schema_without_username(tmp_path) -> None: async def test_migrates_legacy_schema_without_username(tmp_path) -> None:
path = tmp_path / "legacy.db" path = tmp_path / "legacy.db"
connection = sqlite3.connect(path) connection = sqlite3.connect(path)