diff --git a/backend/src/bot/handlers/message/handler.py b/backend/src/bot/handlers/message/handler.py index d44a3bb..f338344 100644 --- a/backend/src/bot/handlers/message/handler.py +++ b/backend/src/bot/handlers/message/handler.py @@ -18,7 +18,7 @@ from aiogram.types import ( ) from convex import ConvexInt64 -from bot.handlers.proxy.handler import get_proxy_config +from bot.handlers.proxy.handler import get_proxy_config, increment_proxy_count from bot.modules.ai import ( SUMMARIZE_PROMPT, AgentDeps, @@ -426,6 +426,8 @@ async def process_message( # noqa: C901, PLR0912, PLR0913, PLR0915 await proxy_config.proxy_bot.send_message( proxy_config.target_chat_id, f"👤 {text}" ) + await increment_proxy_count(chat_id) + proxy_config = get_proxy_config(chat_id) rag_connections = await convex.query( "ragConnections:getActiveForUser", {"userId": convex_user_id} @@ -550,6 +552,7 @@ async def process_message( # noqa: C901, PLR0912, PLR0913, PLR0915 await proxy_config.proxy_bot.send_message( proxy_config.target_chat_id, part ) + await increment_proxy_count(chat_id) except Exception as e: # noqa: BLE001 await state.stop_typing() @@ -663,6 +666,7 @@ async def on_album_message( await proxy_config.proxy_bot.send_media_group( proxy_config.target_chat_id, media ) + await increment_proxy_count(message.chat.id) active_chat_id = user["activeChatId"] await convex.mutation( @@ -732,6 +736,7 @@ async def on_photo_message(message: types.Message, bot: Bot) -> None: BufferedInputFile(image_bytes, "photo.jpg"), caption=f"👤 {caption}", ) + await increment_proxy_count(message.chat.id) active_chat_id = user["activeChatId"] await convex.mutation( diff --git a/backend/src/bot/handlers/proxy/handler.py b/backend/src/bot/handlers/proxy/handler.py index 28086bf..d9e4c7e 100644 --- a/backend/src/bot/handlers/proxy/handler.py +++ b/backend/src/bot/handlers/proxy/handler.py @@ -1,4 +1,5 @@ -from dataclasses import dataclass +import contextlib +from dataclasses import dataclass, field from aiogram import Bot, Router, types from aiogram.client.default import DefaultBotProperties @@ -13,6 +14,8 @@ class ProxyConfig: bot_token: str target_chat_id: int proxy_bot: Bot + limit: int = 0 + message_count: int = field(default=0, init=False) proxy_states: dict[int, ProxyConfig] = {} @@ -28,25 +31,43 @@ async def cleanup_proxy(chat_id: int) -> None: await config.proxy_bot.session.close() +async def increment_proxy_count(chat_id: int) -> bool: + config = proxy_states.get(chat_id) + if not config: + return False + + config.message_count += 1 + + if config.limit > 0 and config.message_count >= config.limit: + await config.proxy_bot.send_message( + config.target_chat_id, f"🔗 Proxy finished ({config.limit} messages)" + ) + await cleanup_proxy(chat_id) + return False + + return True + + @router.message(Command("proxy")) -async def on_proxy(message: types.Message) -> None: # noqa: C901 +async def on_proxy(message: types.Message) -> None: # noqa: C901, PLR0912 if not message.from_user or not message.text: return chat_id = message.chat.id - args = message.text.split(maxsplit=2) + args = message.text.split(maxsplit=3) if len(args) == 1: config = get_proxy_config(chat_id) if config: + cnt = f" ({config.message_count}/{config.limit})" if config.limit else "" await message.answer( - f"Proxy active → chat {config.target_chat_id}\n\n" + f"Proxy active → chat {config.target_chat_id}{cnt}\n\n" "Use /proxy deactivate to stop." ) else: await message.answer( "Usage:\n" - "/proxy BOT_TOKEN CHAT_ID - activate proxy\n" + "/proxy BOT_TOKEN CHAT_ID [LIMIT] - activate proxy\n" "/proxy deactivate - stop proxy" ) return @@ -60,7 +81,7 @@ async def on_proxy(message: types.Message) -> None: # noqa: C901 return if len(args) < 3: # noqa: PLR2004 - await message.answer("Usage: /proxy BOT_TOKEN CHAT_ID") + await message.answer("Usage: /proxy BOT_TOKEN CHAT_ID [LIMIT]") return bot_token = args[1] @@ -70,6 +91,11 @@ async def on_proxy(message: types.Message) -> None: # noqa: C901 await message.answer("Invalid chat ID. Must be a number.") return + limit = 0 + if len(args) >= 4: # noqa: PLR2004 + with contextlib.suppress(ValueError): + limit = max(int(args[3]), 0) + if chat_id in proxy_states: await cleanup_proxy(chat_id) @@ -87,12 +113,16 @@ async def on_proxy(message: types.Message) -> None: # noqa: C901 return proxy_states[chat_id] = ProxyConfig( - bot_token=bot_token, target_chat_id=target_chat_id, proxy_bot=proxy_bot + bot_token=bot_token, + target_chat_id=target_chat_id, + proxy_bot=proxy_bot, + limit=limit, ) + limit_text = f"\nLimit: {limit} messages" if limit > 0 else "" await message.answer( f"✓ Proxy activated via @{bot_info.username}\n" - f"Target: {target_chat_id}\n\n" + f"Target: {target_chat_id}{limit_text}\n\n" "All messages will be forwarded." ) await message.delete()