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

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

README.md

@@ -15,6 +15,7 @@ Telegram-бот для генерации и редактирования изо
 - **Баланс и рефералы**: 3 стартовые генерации, +1 ежедневно в 00:00 МСК, +3 за первую оплату друга.
 - **Управление пользователями**: в веб-интерфейсе можно выдавать админку, менять баланс и блокировать.
 - **Подтверждение 18+**: без подтверждения возраста бот не принимает запросы.
+- **Контекст**: следующий текстовый запрос продолжает предыдущую сцену; `/reset` начинает новую.
 - Глобальная очередь с одной активной задачей на пользователя.
 - Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
 - PostgreSQL: запросы, статусы, seed, пути к файлам.
@@ -201,10 +202,22 @@ TELEGRAM_HEALTH_INTERVAL=60
 - Премиум-очередь всегда оплачивается 15 звёздами; баланс при этом не расходуется.
 - `/balance` — текущий баланс.
 - `/invite` — реферальная ссылка.
+- `/reset` — сбросить контекст и начать новую сцену.
 - За первую оплату приглашённого друга начисляется +3 генерации, реферер получает уведомление.
 
 Баланс и рефералы хранятся в таблице `users`. Пропущенные дни начисляются при следующем запуске/полуночи.
 
+## Контекст
+
+- Первый текстовый запрос создаёт контекст: базовую сцену и текущий улучшенный промпт.
+- Следующий текст обрабатывается вместе с предыдущим промптом: бот сохраняет сцену и применяет только запрошенное изменение.
+- `/reset` очищает контекст; следующий запрос начнёт новую сцену.
+- Новое фото с подписью автоматически начинает новую цепочку.
+- Для точного редактирования конкретной картинки ответь на сообщение бота с изображением.
+- `/reset` не отменяет уже запущенную генерацию: она завершится со своим снимком контекста.
+
+Контекст хранится в таблице `user_contexts` и обновляется только после успешной отправки результата.
+
 ## Оплата
 
 Генерация стоит 5 Telegram Stars в обычной очереди и 15 звёзд в премиум-очереди:

bot.py

@@ -63,7 +63,8 @@ START_TEXT = (
     "Или ответь текстом на готовую картинку — возьму её как исходник.\n\n"
     "/balance — баланс генераций.\n"
     "/invite — пригласить друга.\n"
-    "/status — состояние текущей задачи."
+    "/status — состояние текущей задачи.\n"
+    "/reset — начать новую сцену."
 )
 BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны."
 AGE_TEXT = (
@@ -124,6 +125,7 @@ class MpakoBot:
         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)
         )
@@ -491,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:
@@ -632,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)
@@ -740,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)
@@ -830,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)
@@ -895,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
@@ -914,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 (
@@ -966,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(
                 [
@@ -1023,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(
@@ -1278,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 и веб-история.
-- 59 тестов: 37 проходят локально без БД, 22 теста БД выполняются в CI на PostgreSQL через `TEST_DATABASE_URL`. Покрыты workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, управление пользователями, retry, миграция, watchdog.
+- 62 теста: 38 проходят локально без БД, 24 теста БД выполняются в CI на PostgreSQL через `TEST_DATABASE_URL`. Покрыты workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, управление пользователями, retry, миграция, watchdog, контекст.
 - Не проверено: реальная оплата Stars и реферальная выплата в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. PostgreSQL-миграция и тесты БД ещё не запускались вживую — это сделает CI и первый деплой.
 
 ## 0. Допущения и решения по умолчанию
@@ -381,6 +381,17 @@ ON requests(user_id) WHERE status IN ('queued', 'running');
 
 **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
 
-        workflow = copy.deepcopy(self._template)
+        previous = self._sanitize(previous_prompt or "")
-        workflow[self._text_node_id]["inputs"]["prompt"] = self._chat_prompt(
+        base = self._sanitize(base_prompt or "")
-            user_prompt, mode
+        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"] = chat_prompt
         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)

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)