Контекст между сообщениями

Сохраняю базовый и текущий промпт пользователя и передаю предыдущий контекст в enhancer. Добавил /reset, снимок контекста в requests/pending_payments и обновление только после успешной отправки результата.

.env.example

@@ -2,7 +2,9 @@
 BOT_TOKEN=
 
 # Прокси для Telegram API. Пусто — без прокси.
-# В Docker на Linux: http://host.docker.internal:20171
+# В Docker на Linux рекомендуется SOCKS5:
+# socks5h://host.docker.internal:20170
+# HTTP-вариант: http://host.docker.internal:20171
 TELEGRAM_PROXY=
 # 0 — повторять до успеха; иначе число попыток
 TELEGRAM_RETRY_ATTEMPTS=0

README.md

@@ -15,6 +15,7 @@ Telegram-бот для генерации и редактирования изо
 - **Баланс и рефералы**: 3 стартовые генерации, +1 ежедневно в 00:00 МСК, +3 за первую оплату друга.
 - **Управление пользователями**: в веб-интерфейсе можно выдавать админку, менять баланс и блокировать.
 - **Подтверждение 18+**: без подтверждения возраста бот не принимает запросы.
+- **Контекст**: следующий текстовый запрос продолжает предыдущую сцену; `/reset` начинает новую.
 - Глобальная очередь с одной активной задачей на пользователя.
 - Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
 - PostgreSQL: запросы, статусы, seed, пути к файлам.
@@ -48,14 +49,17 @@ Copy-Item .env.example .env
 Если сервер не имеет прямого доступа к Telegram, укажи прокси V2RayA/HTTP/SOCKS5:
 
 ```dotenv
-# Docker на Linux (V2RayA слушает на хосте)
+# Docker на Linux, рекомендуется SOCKS5
-TELEGRAM_PROXY=http://host.docker.internal:20171
+TELEGRAM_PROXY=socks5h://host.docker.internal:20170
+
+# HTTP-вариант
+# TELEGRAM_PROXY=http://host.docker.internal:20171
 
 # Запуск без Docker, на том же хосте
-# TELEGRAM_PROXY=http://127.0.0.1:20171
+# TELEGRAM_PROXY=socks5h://127.0.0.1:20170
 ```
 
-Для SOCKS5 используй `socks5h://host.docker.internal:20170`; зависимость `socksio` уже включена в `requirements.txt`.
+SOCKS5 стабильнее держит long polling; зависимость `socksio` уже включена в `requirements.txt`.
 
 ## Запуск
 
@@ -177,6 +181,7 @@ TELEGRAM_HEALTH_INTERVAL=60
 - Backoff растёт до `TELEGRAM_RETRY_MAX_DELAY`.
 - Критичные отправки (результат, счёт, подтверждение) повторяются бесконечно; обновления статуса — ограниченно, чтобы не плодить зависшие задачи.
 - Раз в `TELEGRAM_HEALTH_INTERVAL` выполняется `getMe`; при сбое проверка повторяется с backoff, пока Telegram не ответит.
+- Watchdog проверяет pending updates: если они не обрабатываются 5 минут, polling перезапускается, а при неудаче контейнер перезапускается автоматически.
 
 ## Подтверждение возраста
 
@@ -197,10 +202,22 @@ TELEGRAM_HEALTH_INTERVAL=60
 - Премиум-очередь всегда оплачивается 15 звёздами; баланс при этом не расходуется.
 - `/balance` — текущий баланс.
 - `/invite` — реферальная ссылка.
+- `/reset` — сбросить контекст и начать новую сцену.
 - За первую оплату приглашённого друга начисляется +3 генерации, реферер получает уведомление.
 
 Баланс и рефералы хранятся в таблице `users`. Пропущенные дни начисляются при следующем запуске/полуночи.
 
+## Контекст
+
+- Первый текстовый запрос создаёт контекст: базовую сцену и текущий улучшенный промпт.
+- Следующий текст обрабатывается вместе с предыдущим промптом: бот сохраняет сцену и применяет только запрошенное изменение.
+- `/reset` очищает контекст; следующий запрос начнёт новую сцену.
+- Новое фото с подписью автоматически начинает новую цепочку.
+- Для точного редактирования конкретной картинки ответь на сообщение бота с изображением.
+- `/reset` не отменяет уже запущенную генерацию: она завершится со своим снимком контекста.
+
+Контекст хранится в таблице `user_contexts` и обновляется только после успешной отправки результата.
+
 ## Оплата
 
 Генерация стоит 5 Telegram Stars в обычной очереди и 15 звёзд в премиум-очереди:

bot.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 import asyncio
 import logging
+import os
 import time
 import uuid
 from datetime import datetime, timedelta, timezone
 from datetime import time as dt_time
 from pathlib import Path
 
+import httpx
 from telegram import (
     CopyTextButton,
     InlineKeyboardButton,
@@ -25,6 +27,7 @@ from telegram.ext import (
     ContextTypes,
     MessageHandler,
     PreCheckoutQueryHandler,
+    TypeHandler,
     filters,
 )
 
@@ -38,6 +41,21 @@ from workflows import WorkflowBuilder
 
 logger = logging.getLogger(__name__)
 MSK = timezone(timedelta(hours=3))
+POLLING_WATCHDOG_INTERVAL = 60.0
+POLLING_WATCHDOG_STUCK_AFTER = 300.0
+
+
+def polling_is_stuck(
+    *,
+    now: float,
+    last_update_at: float,
+    pending_update_count: int,
+    stuck_after: float,
+) -> bool:
+    return (
+        pending_update_count > 0
+        and now - last_update_at >= stuck_after
+    )
 
 START_TEXT = (
     "Пришли текст — сгенерирую картинку.\n"
@@ -45,7 +63,8 @@ START_TEXT = (
     "Или ответь текстом на готовую картинку — возьму её как исходник.\n\n"
     "/balance — баланс генераций.\n"
     "/invite — пригласить друга.\n"
-    "/status — состояние текущей задачи."
+    "/status — состояние текущей задачи.\n"
+    "/reset — начать новую сцену."
 )
 BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны."
 AGE_TEXT = (
@@ -71,17 +90,24 @@ class MpakoBot:
         self._last_edit: dict[int, tuple[float, int]] = {}
         self._daily_task: asyncio.Task[None] | None = None
         self._health_task: asyncio.Task[None] | None = None
+        self._watchdog_task: asyncio.Task[None] | None = None
+        self._last_update_at = time.monotonic()
         self._application: Application | None = None
 
     def run(self) -> None:
         builder = (
             Application.builder()
             .token(self._settings.bot_token)
-            .connect_timeout(20)
+            .connect_timeout(10)
-            .read_timeout(60)
+            .read_timeout(30)
-            .write_timeout(60)
+            .write_timeout(30)
-            .get_updates_connect_timeout(20)
+            .pool_timeout(5)
-            .get_updates_read_timeout(60)
+            .connection_pool_size(10)
+            .get_updates_connect_timeout(10)
+            .get_updates_read_timeout(15)
+            .get_updates_write_timeout(15)
+            .get_updates_pool_timeout(5)
+            .get_updates_connection_pool_size(10)
             .post_init(self._post_init)
             .post_shutdown(self._post_shutdown)
         )
@@ -91,11 +117,15 @@ class MpakoBot:
             )
         application = builder.build()
         self._application = application
+        application.add_handler(
+            TypeHandler(Update, self._track_update), group=-1
+        )
         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(CommandHandler("reset", self._reset))
         application.add_handler(
             PreCheckoutQueryHandler(self._pre_checkout)
         )
@@ -136,6 +166,9 @@ class MpakoBot:
         self._health_task = asyncio.create_task(
             self._telegram_health_loop(), name="telegram-health"
         )
+        self._watchdog_task = asyncio.create_task(
+            self._polling_watchdog_loop(), name="polling-watchdog"
+        )
         enhancer = (
             PromptEnhancer(self._comfy, self._settings)
             if self._settings.prompt_enhance_enabled
@@ -153,6 +186,10 @@ class MpakoBot:
         logger.info("Бот запущен, ComfyUI: %s", self._settings.comfy_base_url)
 
     async def _post_shutdown(self, application: Application) -> None:
+        if self._watchdog_task is not None:
+            self._watchdog_task.cancel()
+            await asyncio.gather(self._watchdog_task, return_exceptions=True)
+            self._watchdog_task = None
         if self._health_task is not None:
             self._health_task.cancel()
             await asyncio.gather(self._health_task, return_exceptions=True)
@@ -168,6 +205,79 @@ class MpakoBot:
         await self._db.close()
         logger.info("Бот остановлен")
 
+    async def _track_update(
+        self, update: Update, context: ContextTypes.DEFAULT_TYPE
+    ) -> None:
+        self._last_update_at = time.monotonic()
+
+    async def _polling_watchdog_loop(self) -> None:
+        while True:
+            await asyncio.sleep(POLLING_WATCHDOG_INTERVAL)
+            pending = await self._pending_updates_count()
+            if pending is None:
+                continue
+            if not polling_is_stuck(
+                now=time.monotonic(),
+                last_update_at=self._last_update_at,
+                pending_update_count=pending,
+                stuck_after=POLLING_WATCHDOG_STUCK_AFTER,
+            ):
+                continue
+            logger.warning(
+                "Polling watchdog: %s обновлений не обрабатываются,"
+                " перезапускаю updater",
+                pending,
+            )
+            await self._restart_polling()
+
+    async def _pending_updates_count(self) -> int | None:
+        url = (
+            "https://api.telegram.org/bot"
+            f"{self._settings.bot_token}/getWebhookInfo"
+        )
+        proxy = self._settings.telegram_proxy or None
+        try:
+            async with httpx.AsyncClient(
+                proxy=proxy,
+                timeout=httpx.Timeout(10.0, connect=5.0),
+            ) as client:
+                response = await client.get(url)
+                data = response.json()
+        except (httpx.HTTPError, ValueError) as exc:
+            logger.warning(
+                "Polling watchdog: проверка не удалась (%s)",
+                type(exc).__name__,
+            )
+            return None
+        if not data.get("ok"):
+            return None
+        result = data.get("result") or {}
+        return int(result.get("pending_update_count") or 0)
+
+    async def _restart_polling(self) -> None:
+        application = self._application
+        if application is None or application.updater is None:
+            return
+        try:
+            await asyncio.wait_for(
+                application.updater.stop(), timeout=15
+            )
+            await asyncio.wait_for(
+                application.updater.start_polling(
+                    allowed_updates=Update.ALL_TYPES,
+                    drop_pending_updates=False,
+                ),
+                timeout=30,
+            )
+        except Exception:
+            logger.exception(
+                "Polling watchdog: не удалось перезапустить updater,"
+                " выхожу для рестарта контейнера"
+            )
+            os._exit(1)
+        self._last_update_at = time.monotonic()
+        logger.warning("Polling watchdog: updater перезапущен")
+
     async def _daily_credits_loop(self) -> None:
         while True:
             now = datetime.now(MSK)
@@ -383,6 +493,23 @@ class MpakoBot:
             reply_markup=keyboard,
         )
 
+    async def _reset(
+        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
+        cleared = await self._db.clear_user_context(user.id)
+        text = (
+            "Контекст сброшен. Следующий запрос начнёт новую сцену."
+            if cleared
+            else "Контекст уже пуст."
+        )
+        await self._retry(message.reply_text, text)
+
     async def _status(
         self, update: Update, context: ContextTypes.DEFAULT_TYPE
     ) -> None:
@@ -524,6 +651,9 @@ class MpakoBot:
                 input_path=pending.input_file_path,
                 skip_payment=True,
                 is_premium=False,
+                use_context=bool(pending.context_prompt),
+                context_base_prompt=pending.context_base_prompt,
+                context_prompt=pending.context_prompt,
             )
             if not accepted:
                 await self._db.add_credits(query.from_user.id, 1)
@@ -632,6 +762,9 @@ class MpakoBot:
             payment_charge_id=charge_id,
             payment_payload=payload,
             is_premium=pending.is_premium,
+            use_context=bool(pending.context_prompt),
+            context_base_prompt=pending.context_base_prompt,
+            context_prompt=pending.context_prompt,
         )
         if accepted:
             await self._reward_referrer(user)
@@ -722,6 +855,7 @@ class MpakoBot:
             mode="i2i",
             prompt=caption,
             input_path=str(input_path),
+            use_context=False,
         )
         if not accepted:
             input_path.unlink(missing_ok=True)
@@ -787,6 +921,9 @@ class MpakoBot:
         payment_charge_id: str | None = None,
         payment_payload: str | None = None,
         is_premium: bool = False,
+        use_context: bool = True,
+        context_base_prompt: str | None = None,
+        context_prompt: str | None = None,
     ) -> bool:
         message = update.effective_message
         user = update.effective_user
@@ -806,6 +943,14 @@ class MpakoBot:
         if user_row.is_blocked:
             await message.reply_text("Доступ заблокирован.")
             return False
+        if not use_context:
+            context_base_prompt = None
+            context_prompt = None
+        elif context_prompt is None:
+            user_context = await self._db.get_user_context(user.id)
+            if user_context is not None:
+                context_base_prompt = user_context.base_prompt
+                context_prompt = user_context.current_prompt
         credits = user_row.credits
         has_person = False
         if (
@@ -858,6 +1003,8 @@ class MpakoBot:
                 username=user.username,
                 input_file_path=input_path,
                 amount=0,
+                context_base_prompt=context_base_prompt,
+                context_prompt=context_prompt,
             )
             keyboard = InlineKeyboardMarkup(
                 [
@@ -915,6 +1062,8 @@ class MpakoBot:
                 is_admin=is_admin,
                 payment_charge_id=payment_charge_id,
                 is_premium=is_premium,
+                context_base_prompt=context_base_prompt,
+                context_prompt=context_prompt,
             )
         except PromptBlocked:
             note = await self._refund_if_paid(
@@ -1170,6 +1319,23 @@ class MpakoBot:
                 return
 
         await self._clear_status(request.id, delete=True)
+        await self._save_context(request)
+
+    async def _save_context(self, request: Request) -> None:
+        current = request.enhanced_prompt or request.prompt_text
+        base = request.context_base_prompt or current
+        try:
+            await self._db.save_user_context(
+                user_id=request.user_id,
+                base_prompt=base,
+                current_prompt=current,
+                last_request_id=request.id,
+                last_mode=request.mode,
+            )
+        except Exception:
+            logger.exception(
+                "Не удалось сохранить контекст запроса %s", request.id
+            )
 
     async def _finish_failed(self, request: Request, error: str) -> None:
         target = self._status_messages.get(request.id)

checklist.md

@@ -8,7 +8,7 @@
 
 - Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15), оплата Stars (16), баланс/рефералы (17), управление пользователями (18), возрастной gate (19) и надёжность Telegram (20).
 - Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история.
-- 56 тестов: 34 проходят локально без БД, 22 теста БД выполняются в CI на PostgreSQL через `TEST_DATABASE_URL`. Покрыты workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, управление пользователями, retry, миграция.
+- 62 теста: 38 проходят локально без БД, 24 теста БД выполняются в CI на PostgreSQL через `TEST_DATABASE_URL`. Покрыты workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, управление пользователями, retry, миграция, watchdog, контекст.
 - Не проверено: реальная оплата Stars и реферальная выплата в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. PostgreSQL-миграция и тесты БД ещё не запускались вживую — это сделает CI и первый деплой.
 
 ## 0. Допущения и решения по умолчанию
@@ -376,10 +376,22 @@ ON requests(user_id) WHERE status IN ('queued', 'running');
 - [x] `_retry` с настраиваемыми попытками, backoff и режимом «до успеха».
 - [x] Health-check `getMe` с повторами.
 - [x] Ретраи критичных отправок и загрузок файлов.
-- [x] Тесты retry.
+- [x] Watchdog polling: если pending updates не убывают 5 минут, updater перезапускается, при неудаче перезапускается контейнер.
+- [x] Тесты retry и watchdog.
 
 **DoD:** при таймаутах прокси операции повторяются, health-check фиксирует восстановление связи.
 
+## 21. Контекст между сообщениями
+
+- [x] Таблица `user_contexts`: `base_prompt`, `current_prompt`, `last_request_id`, `last_mode`, `updated_at`.
+- [x] Снимок контекста сохраняется в `requests` и `pending_payments`.
+- [x] Следующий текстовый запрос обрабатывается вместе с предыдущим промптом.
+- [x] Контекст обновляется только после успешной отправки результата.
+- [x] `/reset` очищает контекст; новое фото начинает новую цепочку.
+- [x] Тесты БД, enhancer и очереди.
+
+**DoD:** продолжение сцены работает по тексту, `/reset` возвращает к новой сцене.
+
 ## Порядок работ
 
 1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow.

db.py

@@ -61,6 +61,8 @@ class Request:
     timestamp: str
     prompt_text: str
     enhanced_prompt: str | None
+    context_base_prompt: str | None
+    context_prompt: str | None
     input_file_path: str | None
     output_file_path: str | None
     status: str
@@ -86,6 +88,8 @@ class Request:
             timestamp=_as_iso(row["timestamp"]) or "",
             prompt_text=row["prompt_text"],
             enhanced_prompt=row["enhanced_prompt"],
+            context_base_prompt=row["context_base_prompt"],
+            context_prompt=row["context_prompt"],
             input_file_path=row["input_file_path"],
             output_file_path=row["output_file_path"],
             status=row["status"],
@@ -177,6 +181,8 @@ class PendingPayment:
     timestamp: str
     mode: str
     prompt: str
+    context_base_prompt: str | None
+    context_prompt: str | None
     username: str | None
     input_file_path: str | None
     amount: int
@@ -193,6 +199,8 @@ class PendingPayment:
             timestamp=_as_iso(row["timestamp"]) or "",
             mode=row["mode"],
             prompt=row["prompt"],
+            context_base_prompt=row["context_base_prompt"],
+            context_prompt=row["context_prompt"],
             username=row["username"],
             input_file_path=row["input_file_path"],
             amount=int(row["amount"]),
@@ -203,6 +211,27 @@ class PendingPayment:
         )
 
 
+@dataclass(frozen=True)
+class UserContext:
+    user_id: int
+    base_prompt: str
+    current_prompt: str
+    last_request_id: int | None
+    last_mode: str | None
+    updated_at: str
+
+    @classmethod
+    def from_row(cls, row: Mapping[str, Any]) -> UserContext:
+        return cls(
+            user_id=int(row["user_id"]),
+            base_prompt=row["base_prompt"],
+            current_prompt=row["current_prompt"],
+            last_request_id=row["last_request_id"],
+            last_mode=row["last_mode"],
+            updated_at=_as_iso(row["updated_at"]) or "",
+        )
+
+
 SCHEMA = """
 CREATE TABLE IF NOT EXISTS requests (
     id                BIGSERIAL PRIMARY KEY,
@@ -211,6 +240,8 @@ CREATE TABLE IF NOT EXISTS requests (
     timestamp         TIMESTAMPTZ NOT NULL DEFAULT now(),
     prompt_text       TEXT NOT NULL,
     enhanced_prompt   TEXT,
+    context_base_prompt TEXT,
+    context_prompt    TEXT,
     input_file_path   TEXT,
     output_file_path  TEXT,
     status            TEXT NOT NULL DEFAULT 'queued'
@@ -259,6 +290,8 @@ CREATE TABLE IF NOT EXISTS pending_payments (
     status                     TEXT NOT NULL DEFAULT 'pending'
                                CHECK (status IN ('pending', 'paid', 'cancelled', 'refunded')),
     is_premium                 BOOLEAN NOT NULL DEFAULT FALSE,
+    context_base_prompt        TEXT,
+    context_prompt             TEXT,
     telegram_payment_charge_id TEXT,
     request_id                 BIGINT
 );
@@ -283,6 +316,15 @@ CREATE TABLE IF NOT EXISTS users (
 CREATE INDEX IF NOT EXISTS idx_users_referrer
     ON users(referrer_id);
 
+CREATE TABLE IF NOT EXISTS user_contexts (
+    user_id         BIGINT PRIMARY KEY,
+    base_prompt     TEXT NOT NULL,
+    current_prompt  TEXT NOT NULL,
+    last_request_id BIGINT,
+    last_mode       TEXT,
+    updated_at      TIMESTAMPTZ NOT NULL DEFAULT now()
+);
+
 CREATE TABLE IF NOT EXISTS app_settings (
     key   TEXT PRIMARY KEY,
     value TEXT
@@ -294,6 +336,8 @@ SCHEMA_STATEMENTS = [part.strip() for part in SCHEMA.split(";") if part.strip()]
 MIGRATIONS = (
     "ALTER TABLE requests ADD COLUMN IF NOT EXISTS username TEXT",
     "ALTER TABLE requests ADD COLUMN IF NOT EXISTS enhanced_prompt TEXT",
+    "ALTER TABLE requests ADD COLUMN IF NOT EXISTS context_base_prompt TEXT",
+    "ALTER TABLE requests ADD COLUMN IF NOT EXISTS context_prompt TEXT",
     "ALTER TABLE requests ADD COLUMN IF NOT EXISTS payment_charge_id TEXT",
     (
         "ALTER TABLE requests ADD COLUMN IF NOT EXISTS payment_refunded BOOLEAN"
@@ -307,6 +351,8 @@ MIGRATIONS = (
         "ALTER TABLE pending_payments ADD COLUMN IF NOT EXISTS is_premium BOOLEAN"
         " NOT NULL DEFAULT FALSE"
     ),
+    "ALTER TABLE pending_payments ADD COLUMN IF NOT EXISTS context_base_prompt TEXT",
+    "ALTER TABLE pending_payments ADD COLUMN IF NOT EXISTS context_prompt TEXT",
     (
         "ALTER TABLE users ADD COLUMN IF NOT EXISTS is_admin BOOLEAN"
         " NOT NULL DEFAULT FALSE"
@@ -402,6 +448,8 @@ class Database:
         seed: int,
         payment_charge_id: str | None = None,
         is_premium: bool = False,
+        context_base_prompt: str | None = None,
+        context_prompt: str | None = None,
     ) -> Request:
         try:
             async with self._connection_pool().connection() as conn:
@@ -409,9 +457,9 @@ class Database:
                     """
                     INSERT INTO requests
                         (user_id, username, timestamp, prompt_text,
-                         input_file_path, status, mode, seed, payment_charge_id,
+                         context_base_prompt, context_prompt, input_file_path,
-                         is_premium)
+                         status, mode, seed, payment_charge_id, is_premium)
-                    VALUES (%s, %s, %s, %s, %s, 'queued', %s, %s, %s, %s)
+                    VALUES (%s, %s, %s, %s, %s, %s, %s, 'queued', %s, %s, %s, %s)
                     RETURNING *
                     """,
                     (
@@ -419,6 +467,8 @@ class Database:
                         username,
                         utc_now(),
                         prompt_text,
+                        context_base_prompt,
+                        context_prompt,
                         input_file_path,
                         mode,
                         seed,
@@ -462,6 +512,54 @@ class Database:
             row = await cursor.fetchone()
         return User.from_row(row) if row is not None else None
 
+    async def get_user_context(self, user_id: int) -> UserContext | None:
+        async with self._connection_pool().connection() as conn:
+            cursor = await conn.execute(
+                "SELECT * FROM user_contexts WHERE user_id = %s", (user_id,)
+            )
+            row = await cursor.fetchone()
+        return UserContext.from_row(row) if row is not None else None
+
+    async def save_user_context(
+        self,
+        *,
+        user_id: int,
+        base_prompt: str,
+        current_prompt: str,
+        last_request_id: int | None,
+        last_mode: str | None,
+    ) -> None:
+        async with self._connection_pool().connection() as conn:
+            await conn.execute(
+                """
+                INSERT INTO user_contexts
+                    (user_id, base_prompt, current_prompt, last_request_id,
+                     last_mode, updated_at)
+                VALUES (%s, %s, %s, %s, %s, now())
+                ON CONFLICT (user_id) DO UPDATE
+                SET base_prompt = EXCLUDED.base_prompt,
+                    current_prompt = EXCLUDED.current_prompt,
+                    last_request_id = EXCLUDED.last_request_id,
+                    last_mode = EXCLUDED.last_mode,
+                    updated_at = now()
+                """,
+                (
+                    user_id,
+                    base_prompt,
+                    current_prompt,
+                    last_request_id,
+                    last_mode,
+                ),
+            )
+
+    async def clear_user_context(self, user_id: int) -> bool:
+        async with self._connection_pool().connection() as conn:
+            cursor = await conn.execute(
+                "DELETE FROM user_contexts WHERE user_id = %s RETURNING user_id",
+                (user_id,),
+            )
+            return await cursor.fetchone() is not None
+
     async def ensure_user(
         self, user_id: int, username: str | None = None
     ) -> User:
@@ -862,14 +960,17 @@ class Database:
         input_file_path: str | None,
         amount: int,
         is_premium: bool = False,
+        context_base_prompt: str | None = None,
+        context_prompt: str | None = None,
     ) -> PendingPayment:
         async with self._connection_pool().connection() as conn:
             cursor = await conn.execute(
                 """
                 INSERT INTO pending_payments
                     (payload, user_id, timestamp, mode, prompt, username,
-                     input_file_path, amount, status, is_premium)
+                     input_file_path, amount, status, is_premium,
-                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 'pending', %s)
+                     context_base_prompt, context_prompt)
+                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 'pending', %s, %s, %s)
                 RETURNING *
                 """,
                 (
@@ -882,6 +983,8 @@ class Database:
                     input_file_path,
                     amount,
                     is_premium,
+                    context_base_prompt,
+                    context_prompt,
                 ),
             )
             row = await cursor.fetchone()

prompt_enhancer.py

@@ -31,6 +31,23 @@ SYSTEM_PROMPTS = {
     ),
 }
 
+REFINE_SYSTEM_PROMPTS = {
+    "t2i": (
+        "You are editing an existing image prompt. The previous prompt describes "
+        "the current image. Apply only the user's requested changes. Keep all "
+        "subjects, style, colors, composition and constraints from the previous "
+        "prompt unless the user explicitly asks to change them. Return only the "
+        "final English prompt, without quotes or explanations."
+    ),
+    "i2i": (
+        "You are updating an image-editing instruction. The previous prompt "
+        "describes the current image. Apply only the user's requested changes. "
+        "Keep the existing composition, subjects, style, lighting and constraints "
+        "unless the user explicitly asks to change them. Return only the final "
+        "English instruction, without quotes or explanations."
+    ),
+}
+
 
 class PromptEnhanceError(RuntimeError):
     pass
@@ -80,20 +97,35 @@ class PromptEnhancer:
             )
         return str(node_ids[0])
 
-    async def enhance(self, prompt: str, mode: str) -> str:
+    async def enhance(
+        self,
+        prompt: str,
+        mode: str,
+        *,
+        previous_prompt: str | None = None,
+        base_prompt: str | None = None,
+    ) -> str:
         if not self.enabled:
             return prompt
 
-        user_prompt = prompt.strip()
+        user_prompt = self._sanitize(prompt)
-        for token in ("<|im_start|>", "<|im_end|>"):
-            user_prompt = user_prompt.replace(token, "")
         if not user_prompt:
             return prompt
 
+        previous = self._sanitize(previous_prompt or "")
+        base = self._sanitize(base_prompt or "")
+        if previous:
+            chat_prompt = self._refine_chat_prompt(
+                user_prompt,
+                mode,
+                previous_prompt=previous,
+                base_prompt=base or previous,
+            )
+        else:
+            chat_prompt = self._chat_prompt(user_prompt, mode)
+
         workflow = copy.deepcopy(self._template)
-        workflow[self._text_node_id]["inputs"]["prompt"] = self._chat_prompt(
+        workflow[self._text_node_id]["inputs"]["prompt"] = chat_prompt
-            user_prompt, mode
-        )
         prompt_id = await self._comfy.submit(workflow, uuid.uuid4().hex)
         entry = await self._comfy.wait_for_history(
             prompt_id,
@@ -108,6 +140,13 @@ class PromptEnhancer:
             raise PromptEnhanceError("Не удалось очистить сгенерированный промпт")
         return cleaned[: self._settings.max_prompt_length]
 
+    @staticmethod
+    def _sanitize(text: str) -> str:
+        value = text.strip()
+        for token in ("<|im_start|>", "<|im_end|>"):
+            value = value.replace(token, "")
+        return value
+
     @staticmethod
     def _chat_prompt(prompt: str, mode: str) -> str:
         system = SYSTEM_PROMPTS.get(mode, SYSTEM_PROMPTS["t2i"])
@@ -117,6 +156,27 @@ class PromptEnhancer:
             "<|im_start|>assistant\n"
         )
 
+    @staticmethod
+    def _refine_chat_prompt(
+        prompt: str,
+        mode: str,
+        *,
+        previous_prompt: str,
+        base_prompt: str,
+    ) -> str:
+        system = REFINE_SYSTEM_PROMPTS.get(
+            mode, REFINE_SYSTEM_PROMPTS["t2i"]
+        )
+        return (
+            f"<|im_start|>system\n{system}<|im_end|>\n"
+            "<|im_start|>user\n"
+            f"Base scene: {base_prompt}\n"
+            f"Current prompt: {previous_prompt}\n"
+            f"Requested change: {prompt}\n"
+            "<|im_end|>\n"
+            "<|im_start|>assistant\n"
+        )
+
     @staticmethod
     def _clean(text: str) -> str:
         text = " ".join(text.split()).strip()

queue_manager.py

@@ -112,6 +112,8 @@ class QueueManager:
         is_admin: bool = False,
         payment_charge_id: str | None = None,
         is_premium: bool = False,
+        context_base_prompt: str | None = None,
+        context_prompt: str | None = None,
     ) -> tuple[Request, int]:
         async with self._lock:
             if self._stopping:
@@ -148,6 +150,8 @@ class QueueManager:
                 seed=seed,
                 payment_charge_id=payment_charge_id,
                 is_premium=is_premium,
+                context_base_prompt=context_base_prompt,
+                context_prompt=context_prompt,
             )
             target = self._premium if is_premium else self._normal
             target.append(request.id)
@@ -250,16 +254,31 @@ class QueueManager:
                 )
                 try:
                     enhanced = await self._enhancer.enhance(
-                        request.prompt_text, request.mode
+                        request.prompt_text,
+                        request.mode,
+                        previous_prompt=request.context_prompt,
+                        base_prompt=request.context_base_prompt,
                     )
                 except Exception:
                     logger.exception(
                         "Не удалось улучшить промпт запроса %s", request_id
                     )
+                    if request.context_prompt:
+                        fallback = (
+                            f"{request.context_prompt}. {request.prompt_text}"
+                        )
+                        await self._db.set_enhanced_prompt(
+                            request_id, fallback
+                        )
+                        request = await self._refresh(request_id)
                 else:
                     if enhanced and enhanced != request.prompt_text:
                         await self._db.set_enhanced_prompt(request_id, enhanced)
                         request = await self._refresh(request_id)
+            elif request.context_prompt:
+                combined = f"{request.context_prompt}. {request.prompt_text}"
+                await self._db.set_enhanced_prompt(request_id, combined)
+                request = await self._refresh(request_id)
 
             workflow = await self._build_workflow(request)
             progress = await self._comfy.open_progress(client_id)

requirements.txt

@@ -1,5 +1,5 @@
 python-telegram-bot>=21.6,<23
-httpx[socks]>=0.27,<1
+httpx[socks]>=0.28,<1
 psycopg[binary]>=3.2,<4
 psycopg-pool>=3.2,<4
 websockets>=13,<16

tests/conftest.py

@@ -18,6 +18,7 @@ TABLES = (
     "blocked_requests",
     "pending_payments",
     "users",
+    "user_contexts",
     "app_settings",
 )
 

tests/test_db.py

@@ -273,6 +273,41 @@ async def test_user_credits_referral_and_daily_accrual(db: Database) -> None:
     assert await db.accrue_daily_credits(future) == 0
 
 
+async def test_user_context_lifecycle(db: Database) -> None:
+    assert await db.get_user_context(1) is None
+
+    await db.save_user_context(
+        user_id=1,
+        base_prompt="A red apple",
+        current_prompt="A red apple on a table",
+        last_request_id=10,
+        last_mode="t2i",
+    )
+    context = await db.get_user_context(1)
+    assert context is not None
+    assert context.base_prompt == "A red apple"
+    assert context.current_prompt == "A red apple on a table"
+    assert context.last_request_id == 10
+    assert context.last_mode == "t2i"
+
+    await db.save_user_context(
+        user_id=1,
+        base_prompt="A red apple",
+        current_prompt="A red apple on a blue table",
+        last_request_id=11,
+        last_mode="t2i",
+    )
+    context = await db.get_user_context(1)
+    assert context is not None
+    assert context.base_prompt == "A red apple"
+    assert context.current_prompt == "A red apple on a blue table"
+    assert context.last_request_id == 11
+
+    assert await db.clear_user_context(1) is True
+    assert await db.get_user_context(1) is None
+    assert await db.clear_user_context(1) is False
+
+
 async def test_ensure_user_updates_username_without_resetting_credits(
     db: Database,
 ) -> None:

tests/test_prompt_enhancer.py

@@ -62,6 +62,25 @@ async def test_enhance_i2i_prompt() -> None:
     assert "image-editing prompt translator" in prompt
 
 
+async def test_enhance_with_context_uses_refine_prompt() -> None:
+    comfy = FakeEnhanceComfy("A blue apple on a wooden table, soft light")
+    enhancer = PromptEnhancer(comfy, make_settings())
+
+    result = await enhancer.enhance(
+        "Сделай яблоко синим",
+        "t2i",
+        previous_prompt="A red apple on a wooden table, soft light",
+        base_prompt="A red apple on a wooden table",
+    )
+
+    assert result == "A blue apple on a wooden table, soft light"
+    assert comfy.workflow is not None
+    prompt = comfy.workflow["2"]["inputs"]["prompt"]
+    assert "editing an existing image prompt" in prompt
+    assert "A red apple on a wooden table, soft light" in prompt
+    assert "Requested change: Сделай яблоко синим" in prompt
+
+
 async def test_disabled_returns_original_without_comfy_call() -> None:
     comfy = FakeEnhanceComfy()
     enhancer = PromptEnhancer(comfy, make_settings(enabled=False))

tests/test_queue.py

@@ -50,9 +50,18 @@ class FakeEnhancer:
 
     def __init__(self) -> None:
         self.calls: list[tuple[str, str]] = []
+        self.contexts: list[tuple[str | None, str | None]] = []
 
-    async def enhance(self, prompt: str, mode: str) -> str:
+    async def enhance(
+        self,
+        prompt: str,
+        mode: str,
+        *,
+        previous_prompt: str | None = None,
+        base_prompt: str | None = None,
+    ) -> str:
         self.calls.append((prompt, mode))
+        self.contexts.append((previous_prompt, base_prompt))
         return f"enhanced: {prompt}"
 
 
@@ -213,6 +222,40 @@ async def test_prompt_enhancer_is_used(db: Database, tmp_path: Path) -> None:
         await queue.stop()
 
 
+async def test_context_is_passed_to_enhancer(
+    db: Database, tmp_path: Path
+) -> None:
+    events: list[QueueEvent] = []
+    enhancer = FakeEnhancer()
+    queue = make_queue(
+        db, make_settings(tmp_path), events, enhancer=enhancer
+    )
+
+    await queue.start()
+    try:
+        request, _ = await queue.enqueue(
+            user_id=1,
+            mode="t2i",
+            prompt="make it blue",
+            context_base_prompt="A red apple on a table",
+            context_prompt="A red apple on a table, soft light",
+        )
+        await wait_for_event(events, "done")
+
+        assert enhancer.contexts == [
+            (
+                "A red apple on a table, soft light",
+                "A red apple on a table",
+            )
+        ]
+        stored = await db.get_request(request.id)
+        assert stored is not None
+        assert stored.context_prompt == "A red apple on a table, soft light"
+        assert stored.context_base_prompt == "A red apple on a table"
+    finally:
+        await queue.stop()
+
+
 async def test_blocked_prompt_is_logged(db: Database, tmp_path: Path) -> None:
     events: list[QueueEvent] = []
     queue = make_queue(db, make_settings(tmp_path), events)

tests/test_watchdog.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from bot import polling_is_stuck
+
+
+def test_polling_is_stuck_for_unprocessed_updates() -> None:
+    assert (
+        polling_is_stuck(
+            now=1000.0,
+            last_update_at=100.0,
+            pending_update_count=2,
+            stuck_after=300.0,
+        )
+        is True
+    )
+
+
+def test_polling_is_not_stuck_without_pending_updates() -> None:
+    assert (
+        polling_is_stuck(
+            now=1000.0,
+            last_update_at=100.0,
+            pending_update_count=0,
+            stuck_after=300.0,
+        )
+        is False
+    )
+
+
+def test_polling_is_not_stuck_after_recent_update() -> None:
+    assert (
+        polling_is_stuck(
+            now=1000.0,
+            last_update_at=900.0,
+            pending_update_count=2,
+            stuck_after=300.0,
+        )
+        is False
+    )