WEB-APP/ Optimization/ Build fix/ Hotkey edit mode/ Log rotation/ Form a11y/ E2E non-blocking

database/keys.py

@@ -1,9 +1,8 @@
 import asyncio
-from datetime import datetime
+from datetime import UTC, datetime
 from types import SimpleNamespace
 
 from sqlalchemy import delete, func, select, text, update
-from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from core.cache_config import (
@@ -12,6 +11,7 @@ from core.cache_config import (
     KEYS_LIST_CACHE_TTL_SEC,
 )
 from core.redis_cache import cache_delete, cache_get, cache_key, cache_set
+from database.access.resolution import resolve_user_optional
 from database.models import Key, Tariff, User
 from database.users import invalidate_profile_cache, invalidate_user_snapshot
 from logger import logger
@@ -25,10 +25,22 @@ async def invalidate_key_email(client_id: str) -> None:
     await cache_delete(cache_key("key_email", client_id))
 
 
-async def invalidate_keys_list(tg_id: int) -> None:
-    await cache_delete(cache_key("keys_list", tg_id))
-    await cache_delete(cache_key("key_count", tg_id))
-    await invalidate_profile_cache(tg_id)
+async def _purge_keys_cache_ids(*ids: int) -> None:
+    for i in ids:
+        await cache_delete(cache_key("keys_list", i))
+        await cache_delete(cache_key("key_count", i))
+        await invalidate_profile_cache(i)
+
+
+async def invalidate_keys_list(session: AsyncSession, legacy_user_ref: int) -> None:
+    u = await resolve_user_optional(session, legacy_user_ref)
+    if u is None:
+        await _purge_keys_cache_ids(legacy_user_ref)
+        return
+    if u.tg_id is not None:
+        await _purge_keys_cache_ids(u.id, u.tg_id)
+    else:
+        await _purge_keys_cache_ids(u.id)
 
 
 async def invalidate_key_details_by_client_id(session: AsyncSession, client_id: str) -> None:
@@ -45,7 +57,7 @@ async def invalidate_key_details_by_client_id(session: AsyncSession, client_id:
 
 async def store_key(
     session: AsyncSession,
-    tg_id: int,
+    legacy_user_ref: int,
     client_id: str,
     email: str,
     expiry_time: int,
@@ -61,71 +73,72 @@ async def store_key(
     current_traffic_limit: int | None = None,
 ):
     """Сохраняет или обновляет ключ подписки."""
-    try:
-        exists = await session.execute(select(Key).where(Key.tg_id == tg_id, Key.client_id == client_id))
-        existing_key = exists.scalar_one_or_none()
+    u = await resolve_user_optional(session, legacy_user_ref)
+    if u is None:
+        raise ValueError(f"Пользователь не найден для ключа: {legacy_user_ref}")
+    uid = u.id
+    exists = await session.execute(select(Key).where(Key.user_id == uid, Key.client_id == client_id))
+    existing_key = exists.scalar_one_or_none()
 
-        if existing_key:
-            values: dict = {
-                "email": email,
-                "expiry_time": expiry_time,
-                "key": key,
-                "server_id": server_id,
-                "remnawave_link": remnawave_link,
-                "tariff_id": tariff_id,
-                "alias": alias,
-            }
+    if existing_key:
+        values: dict = {
+            "email": email,
+            "expiry_time": expiry_time,
+            "key": key,
+            "server_id": server_id,
+            "remnawave_link": remnawave_link,
+            "tariff_id": tariff_id,
+            "alias": alias,
+            "tg_id": u.tg_id,
+        }
 
-            if selected_device_limit is not None:
-                values["selected_device_limit"] = selected_device_limit
-            if selected_traffic_limit is not None:
-                values["selected_traffic_limit"] = selected_traffic_limit
-            if selected_price_rub is not None:
-                values["selected_price_rub"] = selected_price_rub
-            if current_device_limit is not None:
-                values["current_device_limit"] = current_device_limit
-            if current_traffic_limit is not None:
-                values["current_traffic_limit"] = current_traffic_limit
+        if selected_device_limit is not None:
+            values["selected_device_limit"] = selected_device_limit
+        if selected_traffic_limit is not None:
+            values["selected_traffic_limit"] = selected_traffic_limit
+        if selected_price_rub is not None:
+            values["selected_price_rub"] = selected_price_rub
+        if current_device_limit is not None:
+            values["current_device_limit"] = current_device_limit
+        if current_traffic_limit is not None:
+            values["current_traffic_limit"] = current_traffic_limit
 
-            await session.execute(update(Key).where(Key.tg_id == tg_id, Key.client_id == client_id).values(**values))
-            logger.info(f"[Store Key] Ключ обновлён: tg_id={tg_id}, client_id={client_id}, server_id={server_id}")
-        else:
-            if current_device_limit is None:
-                current_device_limit = selected_device_limit
-            if current_traffic_limit is None:
-                current_traffic_limit = selected_traffic_limit
+        await session.execute(update(Key).where(Key.user_id == uid, Key.client_id == client_id).values(**values))
+        logger.info(f"[Store Key] Ключ обновлён: user_id={uid}, client_id={client_id}, server_id={server_id}")
+    else:
+        if current_device_limit is None:
+            current_device_limit = selected_device_limit
+        if current_traffic_limit is None:
+            current_traffic_limit = selected_traffic_limit
 
-            new_key = Key(
-                tg_id=tg_id,
-                client_id=client_id,
-                email=email,
-                created_at=int(datetime.utcnow().timestamp() * 1000),
-                expiry_time=expiry_time,
-                key=key,
-                server_id=server_id,
-                remnawave_link=remnawave_link,
-                tariff_id=tariff_id,
-                alias=alias,
-                selected_device_limit=selected_device_limit,
-                selected_traffic_limit=selected_traffic_limit,
-                selected_price_rub=selected_price_rub,
-                current_device_limit=current_device_limit,
-                current_traffic_limit=current_traffic_limit,
-            )
-            add_result = session.add(new_key)
-            if asyncio.iscoroutine(add_result):
-                await add_result
-            logger.info(f"[Store Key] Ключ создан: tg_id={tg_id}, client_id={client_id}, server_id={server_id}")
+        new_key = Key(
+            user_id=uid,
+            tg_id=u.tg_id,
+            client_id=client_id,
+            email=email,
+            created_at=int(datetime.now(UTC).timestamp() * 1000),
+            expiry_time=expiry_time,
+            key=key,
+            server_id=server_id,
+            remnawave_link=remnawave_link,
+            tariff_id=tariff_id,
+            alias=alias,
+            selected_device_limit=selected_device_limit,
+            selected_traffic_limit=selected_traffic_limit,
+            selected_price_rub=selected_price_rub,
+            current_device_limit=current_device_limit,
+            current_traffic_limit=current_traffic_limit,
+        )
+        add_result = session.add(new_key)
+        if asyncio.iscoroutine(add_result):
+            await add_result
+        logger.info(f"[Store Key] Ключ создан: user_id={uid}, client_id={client_id}, server_id={server_id}")
 
-        await session.commit()
-        invalidate_user_snapshot(tg_id)
-        await invalidate_keys_list(tg_id)
-        await invalidate_key_details(email)
-
-    except SQLAlchemyError as e:
-        logger.error(f"❌ Ошибка при сохранении ключа: {e}")
-        await session.rollback()
-        raise
+    invalidate_user_snapshot(uid)
+    if u.tg_id is not None:
+        invalidate_user_snapshot(u.tg_id)
+    await invalidate_keys_list(session, uid)
+    await invalidate_key_details(email)
 
 
 def _key_to_cache_dict(k: Key) -> dict:
@@ -143,12 +156,16 @@ def _key_to_cache_dict(k: Key) -> dict:
     }
 
 
-async def get_keys(session: AsyncSession, tg_id: int):
-    ckey = cache_key("keys_list", tg_id)
+async def get_keys(session: AsyncSession, legacy_user_ref: int):
+    u = await resolve_user_optional(session, legacy_user_ref)
+    if u is None:
+        return []
+    uid = u.id
+    ckey = cache_key("keys_list", uid)
     cached = await cache_get(ckey)
     if isinstance(cached, list):
         return [SimpleNamespace(**d) for d in cached]
-    result = await session.execute(select(Key).where(Key.tg_id == tg_id))
+    result = await session.execute(select(Key).where(Key.user_id == uid))
     rows = result.scalars().all()
     serialized = [_key_to_cache_dict(k) for k in rows]
     await cache_set(ckey, serialized, KEYS_LIST_CACHE_TTL_SEC)
@@ -160,24 +177,33 @@ async def get_all_keys(session: AsyncSession):
     return result.scalars().all()
 
 
-async def get_key_by_server(session: AsyncSession, tg_id: int, client_id: str):
-    stmt = select(Key).where(Key.tg_id == tg_id, Key.client_id == client_id)
+async def get_key_by_server(session: AsyncSession, legacy_user_ref: int, client_id: str):
+    u = await resolve_user_optional(session, legacy_user_ref)
+    if u is None:
+        return None
+    stmt = select(Key).where(Key.user_id == u.id, Key.client_id == client_id)
     result = await session.execute(stmt)
     return result.scalar_one_or_none()
 
 
-async def get_key_by_email(session: AsyncSession, email: str, tg_id: int | None = None) -> Key | None:
+async def get_key_by_email(session: AsyncSession, email: str, legacy_user_ref: int | None = None) -> Key | None:
     stmt = select(Key).where(Key.email == email)
-    if tg_id is not None:
-        stmt = stmt.where(Key.tg_id == tg_id)
+    if legacy_user_ref is not None:
+        u = await resolve_user_optional(session, legacy_user_ref)
+        if u is None:
+            return None
+        stmt = stmt.where(Key.user_id == u.id)
     result = await session.execute(stmt.limit(1))
     return result.scalar_one_or_none()
 
 
-async def get_key_by_client_id(session: AsyncSession, client_id: str, tg_id: int | None = None) -> Key | None:
+async def get_key_by_client_id(session: AsyncSession, client_id: str, legacy_user_ref: int | None = None) -> Key | None:
     stmt = select(Key).where(Key.client_id == client_id)
-    if tg_id is not None:
-        stmt = stmt.where(Key.tg_id == tg_id)
+    if legacy_user_ref is not None:
+        u = await resolve_user_optional(session, legacy_user_ref)
+        if u is None:
+            return None
+        stmt = stmt.where(Key.user_id == u.id)
     result = await session.execute(stmt.limit(1))
     return result.scalar_one_or_none()
 
@@ -218,15 +244,15 @@ async def get_key_details(session: AsyncSession, email: str) -> dict | None:
     if isinstance(cached, dict):
         return cached
 
-    stmt = select(Key, User).join(User, Key.tg_id == User.tg_id).where(Key.email == email)
+    stmt = select(Key, User).join(User, Key.user_id == User.id).where(Key.email == email)
     result = await session.execute(stmt)
     row = result.first()
     if not row:
         return None
 
     key, user = row
-    expiry_date = datetime.utcfromtimestamp(key.expiry_time / 1000)
-    current_date = datetime.utcnow()
+    expiry_date = datetime.fromtimestamp(key.expiry_time / 1000, UTC)
+    current_date = datetime.now(UTC)
     time_left = expiry_date - current_date
 
     if time_left.total_seconds() <= 0:
@@ -267,39 +293,155 @@ async def get_key_details(session: AsyncSession, email: str) -> dict | None:
     return out
 
 
-async def get_key_count(session: AsyncSession, tg_id: int) -> int:
-    cached = await cache_get(cache_key("key_count", tg_id))
+async def get_key_count(session: AsyncSession, legacy_user_ref: int) -> int:
+    u = await resolve_user_optional(session, legacy_user_ref)
+    if u is None:
+        return 0
+    uid = u.id
+    cached = await cache_get(cache_key("key_count", uid))
     if cached is not None:
         try:
             return int(cached)
         except (TypeError, ValueError):
             pass
-    result = await session.execute(select(func.count()).select_from(Key).where(Key.tg_id == tg_id))
+    result = await session.execute(select(func.count()).select_from(Key).where(Key.user_id == uid))
     count = result.scalar() or 0
-    await cache_set(cache_key("key_count", tg_id), count, KEY_COUNT_CACHE_TTL_SEC)
+    await cache_set(cache_key("key_count", uid), count, KEY_COUNT_CACHE_TTL_SEC)
     return count
 
 
-async def delete_key(session: AsyncSession, identifier: int | str, commit: bool = True):
-    tg_id_for_cache = None
+async def get_key_by_user_and_email(session: AsyncSession, user_id: int, email: str) -> Key | None:
+    """Возвращает ORM-объект Key по паре (users.id, email) или None."""
+    result = await session.execute(
+        select(Key).where(Key.user_id == int(user_id), Key.email == email)
+    )
+    return result.scalar_one_or_none()
+
+
+async def delete_key_by_user_and_email(session: AsyncSession, user_id: int, email: str) -> None:
+    """Удаляет ключ по паре (users.id, email). Commit — ответственность caller'а."""
+    await session.execute(
+        delete(Key).where(Key.user_id == int(user_id), Key.email == email)
+    )
+
+
+async def get_user_keys_with_servers_by_email(
+    session: AsyncSession, user_id: int, email: str
+) -> list[tuple[str, str, dict]]:
+    """Возвращает ключи пользователя + инфо о серверах (join Key × Server).
+
+    Каждый элемент — ``(client_id, server_id, server_info_dict)``. Join
+    делается по (Key.server_id == Server.server_name OR Server.cluster_name),
+    чтобы поддержать и country-mode (server_id = cluster), и cluster-mode
+    (server_id = server_name).
+
+    Используется в ``services.operations.traffic.get_user_traffic``.
+    """
+    from sqlalchemy import or_
+
+    from database.models import Server
+
+    join_cond = or_(
+        Key.server_id == Server.server_name,
+        Key.server_id == Server.cluster_name,
+    )
+    result = await session.execute(
+        select(Key.client_id, Key.server_id, Server)
+        .select_from(Key)
+        .join(Server, join_cond)
+        .where(Server.enabled.is_(True), Key.user_id == int(user_id), Key.email == email)
+    )
+    rows = []
+    for client_id, server_id, server in result.all():
+        rows.append((
+            client_id,
+            server_id,
+            {
+                "server_name": server.server_name,
+                "cluster_name": server.cluster_name,
+                "api_url": server.api_url,
+                "panel_type": server.panel_type,
+            },
+        ))
+    return rows
+
+
+async def get_key_client_id_by_email_and_server(
+    session: AsyncSession, email: str, server_id: str
+) -> str | None:
+    """Возвращает ``client_id`` первого ключа для пары (email, server_id).
+
+    Используется для remnawave traffic reset, где нам нужен только client_id,
+    без остальных полей ключа.
+    """
+    result = await session.execute(
+        select(Key.client_id)
+        .where(Key.email == email, Key.server_id == server_id)
+        .limit(1)
+    )
+    return result.scalar()
+
+
+async def count_keys_by_server_id(session: AsyncSession, server_id: str) -> int:
+    """Сколько всего ключей привязано к указанному server_id (кластеру или серверу).
+
+    Используется для проверки max_keys лимита. ``server_id`` — строка
+    (у ``keys.server_id`` колонка типа String, содержит либо cluster_name,
+    либо server_name в зависимости от страны/кластера).
+    """
+    result = await session.execute(
+        select(func.count()).select_from(Key).where(Key.server_id == server_id)
+    )
+    return int(result.scalar() or 0)
+
+
+async def get_all_key_server_ids(session: AsyncSession) -> list[str]:
+    """Список всех ``server_id`` из таблицы keys (с повторениями).
+
+    Используется в ``services.clusters.select_cluster`` для подсчёта загрузки
+    кластеров. Возвращаем только server_id строки без подгрузки остальных
+    полей, чтобы не тянуть сотни мегабайт для огромных deployments.
+    """
+    result = await session.execute(select(Key.server_id))
+    return [row[0] for row in result.all() if row[0] is not None]
+
+
+async def count_active_keys_for_user(session: AsyncSession, user_id: int) -> int:
+    """Количество незамороженных ключей у пользователя (по internal users.id).
+
+    Отличается от `get_key_count`: не кэшируется и явно исключает замороженные.
+    Используется в проверке "новый пользователь" для купонных правил.
+    """
+    result = await session.execute(
+        select(func.count())
+        .select_from(Key)
+        .where(Key.user_id == int(user_id), Key.is_frozen.is_(False))
+    )
+    return int(result.scalar() or 0)
+
+
+async def delete_key(session: AsyncSession, identifier: int | str):
+    legacy_for_cache = None
     email_for_cache = None
     if isinstance(identifier, str):
         res = await session.execute(
-            select(Key.tg_id, Key.email).where(Key.client_id == identifier).limit(1)
+            select(Key.user_id, Key.email).where(Key.client_id == identifier).limit(1)
         )
         row = res.first()
         if row:
-            tg_id_for_cache, email_for_cache = row[0], row[1]
+            legacy_for_cache, email_for_cache = row[0], row[1]
         await cache_delete(cache_key("key_email", identifier))
+        await session.execute(delete(Key).where(Key.client_id == identifier))
     else:
-        tg_id_for_cache = identifier
-    stmt = delete(Key).where(Key.tg_id == identifier if isinstance(identifier, int) else Key.client_id == identifier)
-    await session.execute(stmt)
-    if commit:
-        await session.commit()
-    if tg_id_for_cache is not None:
-        invalidate_user_snapshot(tg_id_for_cache)
-        await invalidate_keys_list(tg_id_for_cache)
+        u = await resolve_user_optional(session, identifier)
+        if u is None:
+            logger.info(f"Ключ не удалён: пользователь {identifier} не найден")
+            return
+        legacy_for_cache = u.id
+        await session.execute(delete(Key).where(Key.user_id == u.id))
+    if legacy_for_cache is not None:
+        invalidate_user_snapshot(legacy_for_cache)
+        await invalidate_keys_list(session, legacy_for_cache)
     if email_for_cache is not None:
         await invalidate_key_details(str(email_for_cache))
     logger.info(f"Ключ с идентификатором {identifier} удалён")
@@ -307,7 +449,6 @@ async def delete_key(session: AsyncSession, identifier: int | str, commit: bool
 
 async def update_key_expiry(session: AsyncSession, client_id: str, new_expiry_time: int):
     await session.execute(update(Key).where(Key.client_id == client_id).values(expiry_time=new_expiry_time))
-    await session.commit()
     await invalidate_key_details_by_client_id(session, client_id)
     logger.info(f"Срок действия ключа {client_id} обновлён до {new_expiry_time}")
 
@@ -317,59 +458,120 @@ async def get_client_id_by_email(session: AsyncSession, email: str):
     return result.scalar_one_or_none()
 
 
-async def update_key_notified(session: AsyncSession, tg_id: int, client_id: str):
-    await session.execute(update(Key).where(Key.tg_id == tg_id, Key.client_id == client_id).values(notified=True))
-    await session.commit()
-    await invalidate_keys_list(tg_id)
+async def update_key_notified(session: AsyncSession, legacy_user_ref: int, client_id: str):
+    u = await resolve_user_optional(session, legacy_user_ref)
+    if u is None:
+        return
+    await session.execute(update(Key).where(Key.user_id == u.id, Key.client_id == client_id).values(notified=True))
+    await invalidate_keys_list(session, u.id)
     await invalidate_key_details_by_client_id(session, client_id)
 
 
-async def mark_key_as_frozen(session: AsyncSession, tg_id: int, client_id: str, time_left: int):
+async def mark_key_as_frozen(session: AsyncSession, legacy_user_ref: int, client_id: str, time_left: int):
+    u = await resolve_user_optional(session, legacy_user_ref)
+    if u is None:
+        return
     await session.execute(
         text(
             """
             UPDATE keys
             SET expiry_time = :expiry,
                 is_frozen = TRUE
-            WHERE tg_id = :tg_id
+            WHERE user_id = :user_id
               AND client_id = :client_id
             """
         ),
-        {"expiry": time_left, "tg_id": tg_id, "client_id": client_id},
+        {"expiry": time_left, "user_id": u.id, "client_id": client_id},
     )
-    await invalidate_keys_list(tg_id)
+    await invalidate_keys_list(session, u.id)
     await invalidate_key_details_by_client_id(session, client_id)
 
 
 async def mark_key_as_unfrozen(
     session: AsyncSession,
-    tg_id: int,
+    legacy_user_ref: int,
     client_id: str,
     new_expiry_time: int,
 ):
+    u = await resolve_user_optional(session, legacy_user_ref)
+    if u is None:
+        return
     await session.execute(
         text(
             """
             UPDATE keys
             SET expiry_time = :expiry,
                 is_frozen = FALSE
-            WHERE tg_id = :tg_id
+            WHERE user_id = :user_id
               AND client_id = :client_id
             """
         ),
-        {"expiry": new_expiry_time, "tg_id": tg_id, "client_id": client_id},
+        {"expiry": new_expiry_time, "user_id": u.id, "client_id": client_id},
     )
-    await invalidate_keys_list(tg_id)
+    await invalidate_keys_list(session, u.id)
     await invalidate_key_details_by_client_id(session, client_id)
 
 
 async def update_key_tariff(session: AsyncSession, client_id: str, tariff_id: int):
     await session.execute(update(Key).where(Key.client_id == client_id).values(tariff_id=tariff_id))
-    await session.commit()
     await invalidate_key_details_by_client_id(session, client_id)
     logger.info(f"Тариф ключа {client_id} обновлён на {tariff_id}")
 
 
+async def update_key_renewal_snapshot(
+    session: AsyncSession,
+    email: str,
+    *,
+    tariff_id: int,
+    selected_device_limit: int | None = None,
+    current_device_limit: int | None = None,
+    selected_traffic_limit: int | None = None,
+    current_traffic_limit: int | None = None,
+    apply_limits: bool = True,
+) -> None:
+    """Обновляет tariff_id и (опционально) лимиты ключа после продления.
+
+    ``apply_limits=True`` — выставить все четыре лимита (для non-configurable
+    тарифов). ``apply_limits=False`` — обновить только ``tariff_id``, лимиты
+    не трогать (configurable-тарифы обновляют их через `save_key_config_with_mode`).
+    """
+    values: dict = {"tariff_id": tariff_id}
+    if apply_limits:
+        values["selected_device_limit"] = selected_device_limit
+        values["current_device_limit"] = current_device_limit
+        values["selected_traffic_limit"] = selected_traffic_limit
+        values["current_traffic_limit"] = current_traffic_limit
+    await session.execute(update(Key).where(Key.email == email).values(**values))
+    await invalidate_key_details(email)
+
+
+async def update_key_post_creation_snapshot(
+    session: AsyncSession,
+    *,
+    user_id: int,
+    email: str,
+    selected_device_limit: int | None,
+    selected_traffic_limit: int | None,
+    selected_price_rub: int | None,
+) -> None:
+    """Дозаписывает выбранные пользователем параметры ключа сразу после создания.
+
+    Используется из `services.keys.create_vpn_key_headless` — тариф/лимиты не
+    всегда известны на момент `create_key_on_cluster`, поэтому после него
+    идёт snapshot-апдейт для полей, которые нужны для отображения в UI.
+    """
+    await session.execute(
+        update(Key)
+        .where(Key.user_id == int(user_id), Key.email == email)
+        .values(
+            selected_device_limit=selected_device_limit,
+            selected_traffic_limit=selected_traffic_limit,
+            selected_price_rub=selected_price_rub,
+        )
+    )
+    await invalidate_key_details(email)
+
+
 async def get_subscription_link(session: AsyncSession, email: str) -> str | None:
     result = await session.execute(select(func.coalesce(Key.key, Key.remnawave_link)).where(Key.email == email))
     return result.scalar_one_or_none()
@@ -377,7 +579,6 @@ async def get_subscription_link(session: AsyncSession, email: str) -> str | None
 
 async def update_key_client_id(session: AsyncSession, email: str, new_client_id: str):
     await session.execute(update(Key).where(Key.email == email).values(client_id=new_client_id))
-    await session.commit()
     await invalidate_key_details(email)
     logger.info(f"client_id обновлён для {email} -> {new_client_id}")
 
@@ -385,7 +586,6 @@ async def update_key_client_id(session: AsyncSession, email: str, new_client_id:
 async def update_key_link(session: AsyncSession, email: str, link: str) -> bool:
     q = update(Key).where(Key.email == email).values(key=link).returning(Key.client_id)
     res = await session.execute(q)
-    await session.commit()
     ok = res.scalar_one_or_none() is not None
     if ok:
         await invalidate_key_details(email)
@@ -403,7 +603,6 @@ async def update_key_subscription_links(session: AsyncSession, email: str, link:
         .returning(Key.client_id)
     )
     res = await session.execute(stmt)
-    await session.commit()
     ok = res.scalar_one_or_none() is not None
     if ok:
         await invalidate_key_details(email)
@@ -444,10 +643,13 @@ async def save_key_config_with_mode(
     await invalidate_key_details(email)
 
 
-async def reset_key_tariff_state(session: AsyncSession, tg_id: int, email: str, tariff_id: int) -> None:
+async def reset_key_tariff_state(session: AsyncSession, legacy_user_ref: int, email: str, tariff_id: int) -> None:
+    u = await resolve_user_optional(session, legacy_user_ref)
+    if u is None:
+        return
     await session.execute(
         update(Key)
-        .where(Key.tg_id == tg_id, Key.email == email)
+        .where(Key.user_id == u.id, Key.email == email)
         .values(
             tariff_id=tariff_id,
             selected_device_limit=None,
@@ -457,25 +659,27 @@ async def reset_key_tariff_state(session: AsyncSession, tg_id: int, email: str,
             selected_price_rub=None,
         )
     )
-    await session.commit()
-    await invalidate_keys_list(tg_id)
+    await invalidate_keys_list(session, u.id)
     await invalidate_key_details(email)
 
 
 async def save_key_tariff_selection(
     session: AsyncSession,
-    tg_id: int,
+    legacy_user_ref: int,
     email: str,
     tariff_id: int,
     selected_devices: int | None,
     selected_traffic_gb: int | None,
 ) -> None:
+    u = await resolve_user_optional(session, legacy_user_ref)
+    if u is None:
+        return
     selected_devices_val = int(selected_devices) if selected_devices is not None else None
     selected_traffic_val = int(selected_traffic_gb) if selected_traffic_gb is not None and int(selected_traffic_gb) > 0 else None
 
     await session.execute(
         update(Key)
-        .where(Key.tg_id == tg_id, Key.email == email)
+        .where(Key.user_id == u.id, Key.email == email)
         .values(
             tariff_id=tariff_id,
             selected_device_limit=selected_devices_val,
@@ -485,8 +689,7 @@ async def save_key_tariff_selection(
             selected_price_rub=None,
         )
     )
-    await session.commit()
-    await invalidate_keys_list(tg_id)
+    await invalidate_keys_list(session, u.id)
     await invalidate_key_details(email)
 
 
@@ -510,7 +713,6 @@ async def save_admin_key_config(
             selected_price_rub=selected_price,
         )
     )
-    await session.commit()
     await invalidate_key_details(email)
 
 
@@ -527,6 +729,5 @@ async def reset_key_current_limits_to_selected(session: AsyncSession, client_id:
         ),
         {"client_id": client_id},
     )
-    await session.commit()
     await invalidate_key_details_by_client_id(session, client_id)
     logger.info(f"Текущие лимиты ключа {client_id} сброшены к выбранным")