Add multi-user AGY authentication, OAuth login and profile isolation

telegram-agy-bot/agy_engine.py

@@ -13,6 +13,7 @@ from typing import Optional, Dict, Any, Callable, List, Tuple
 from dataclasses import dataclass, field, asdict
 
 from config import settings
+from auth_manager import auth_manager
 
 logger = logging.getLogger(__name__)
 
@@ -1625,6 +1626,7 @@ class AGYEngine:
             except Exception as pe:
                 logger.debug(f"Pre-prompt git pull skipped: {pe}")
 
+        user_env = auth_manager.get_user_env(session.chat_id)
         start_time = time.time()
         proc = await asyncio.create_subprocess_exec(
             *cmd,
@@ -1632,6 +1634,7 @@ class AGYEngine:
             stderr=asyncio.subprocess.PIPE,
             limit=100 * 1024 * 1024,  # 100 MB buffer limit to handle large stream-json chunks
             cwd=workspace_path,
+            env=user_env,
             start_new_session=True,
         )
 

telegram-agy-bot/auth_manager.py

@@ -0,0 +1,309 @@
+import os
+import re
+import json
+import time
+import shutil
+import asyncio
+import logging
+from pathlib import Path
+from typing import Optional, Dict, Any, Tuple
+from datetime import datetime
+
+logger = logging.getLogger("AGYAuthManager")
+
+PROFILES_BASE_DIR = Path("/root/.gemini_profiles")
+MASTER_TOKEN_PATH = Path("/root/.gemini/antigravity-cli/antigravity-oauth-token")
+MASTER_SETTINGS_PATH = Path("/root/.gemini/antigravity-cli/settings.json")
+
+def escape_pre(text: str) -> str:
+    """Escapes HTML entities for pre block."""
+    if not text:
+        return ""
+    return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
+
+class AuthManager:
+    """
+    Manages multi-account authentication, isolated user profiles,
+    and OAuth login flows for Antigravity (AGY) CLI.
+    """
+
+    def __init__(self, base_dir: Path = PROFILES_BASE_DIR):
+        self.base_dir = base_dir
+        self.base_dir.mkdir(parents=True, exist_ok=True)
+        # Stores active login subprocesses: {user_id: {"proc": Process, "url": str, "started_at": float}}
+        self.active_logins: Dict[int, Dict[str, Any]] = {}
+
+    def get_profile_dir(self, user_id: int) -> Path:
+        """Returns the isolated home directory for a given user."""
+        return self.base_dir / str(user_id)
+
+    def get_config_dir(self, user_id: int) -> Path:
+        """Returns the .gemini/antigravity-cli directory inside user profile."""
+        return self.get_profile_dir(user_id) / ".gemini" / "antigravity-cli"
+
+    def get_token_path(self, user_id: int) -> Path:
+        """Returns the path to the user's antigravity-oauth-token."""
+        return self.get_config_dir(user_id) / "antigravity-oauth-token"
+
+    def has_custom_account(self, user_id: int) -> bool:
+        """Checks if a user has a valid personal AGY account token configured."""
+        token_file = self.get_token_path(user_id)
+        if token_file.exists() and token_file.is_file() and token_file.stat().st_size > 10:
+            return True
+        return False
+
+    def get_user_home(self, user_id: int) -> str:
+        """Returns the HOME directory path to use for this user."""
+        if self.has_custom_account(user_id):
+            return str(self.get_profile_dir(user_id))
+        return "/root"
+
+    def get_user_env(self, user_id: int) -> Dict[str, str]:
+        """Returns environment variables dictionary with the proper HOME set for the user."""
+        env = os.environ.copy()
+        env["HOME"] = self.get_user_home(user_id)
+        return env
+
+    def get_account_status(self, user_id: int) -> Dict[str, Any]:
+        """Returns detailed status of the user's active account."""
+        has_custom = self.has_custom_account(user_id)
+        token_path = self.get_token_path(user_id) if has_custom else MASTER_TOKEN_PATH
+
+        modified_time_str = "نامشخص"
+        if token_path.exists():
+            try:
+                mtime = token_path.stat().st_mtime
+                modified_time_str = datetime.fromtimestamp(mtime).strftime("%Y-%m-%d %H:%M:%S")
+            except Exception:
+                pass
+
+        return {
+            "user_id": user_id,
+            "has_custom_account": has_custom,
+            "is_default_server": not has_custom,
+            "account_type": "حساب اختصاصی (Custom Account)" if has_custom else "حساب پیش‌فرض سرور (Server Default)",
+            "token_path": str(token_path),
+            "last_modified": modified_time_str,
+            "is_login_in_progress": user_id in self.active_logins,
+        }
+
+    async def start_oauth_login(self, user_id: int, timeout: int = 25) -> Tuple[bool, str, Optional[str]]:
+        """
+        Starts headless OAuth login flow for a user:
+        Launches AGY CLI in a clean/temporary profile state, intercepts the OAuth URL,
+        and leaves the subprocess waiting for authorization code input.
+        """
+        # Cancel any previous login session for this user
+        self.cancel_oauth_login(user_id)
+
+        config_dir = self.get_config_dir(user_id)
+        config_dir.mkdir(parents=True, exist_ok=True)
+
+        # Temporary backup existing token if user wants to re-login
+        token_file = self.get_token_path(user_id)
+        backup_token = None
+        if token_file.exists():
+            backup_token = token_file.read_text(encoding="utf-8", errors="ignore")
+            try:
+                token_file.unlink()
+            except Exception:
+                pass
+
+        env = os.environ.copy()
+        env["HOME"] = str(self.get_profile_dir(user_id))
+
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                "agy", "--print", "/usage",
+                stdin=asyncio.subprocess.PIPE,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                env=env,
+                cwd="/root",
+                start_new_session=True,
+            )
+
+            auth_url: Optional[str] = None
+            start_time = time.time()
+            url_pattern = re.compile(r"https://accounts\.google\.com/o/oauth2/auth\S+")
+
+            while time.time() - start_time < timeout:
+                try:
+                    line_bytes = await asyncio.wait_for(proc.stdout.readline(), timeout=3.0)
+                    if not line_bytes:
+                        break
+                    line_str = line_bytes.decode("utf-8", errors="replace").strip()
+                    logger.debug(f"OAuth login stdout line for user {user_id}: {line_str}")
+
+                    match = url_pattern.search(line_str)
+                    if match:
+                        auth_url = match.group(0).rstrip(").,")
+                        break
+
+                    if "Waiting for authentication" in line_str or "paste the authorization code" in line_str:
+                        if auth_url:
+                            break
+                except asyncio.TimeoutError:
+                    if proc.returncode is not None:
+                        break
+                    continue
+
+            if not auth_url:
+                try:
+                    proc.kill()
+                except Exception:
+                    pass
+                if backup_token:
+                    token_file.write_text(backup_token, encoding="utf-8")
+                    token_file.chmod(0o600)
+                return False, "عدم دریافت لینک احراز هویت از سرور گوگل. لطفاً مجدداً تلاش فرمایید.", None
+
+            self.active_logins[user_id] = {
+                "proc": proc,
+                "url": auth_url,
+                "backup_token": backup_token,
+                "started_at": time.time(),
+            }
+
+            return True, "لینک احراز هویت گوگل با موفقیت ایجاد شد.", auth_url
+
+        except Exception as e:
+            logger.error(f"Error starting OAuth login for user {user_id}: {e}", exc_info=True)
+            if backup_token:
+                try:
+                    token_file.write_text(backup_token, encoding="utf-8")
+                    token_file.chmod(0o600)
+                except Exception:
+                    pass
+            return False, f"خطا در ایجاد فرآیند لاگین: {e}", None
+
+    async def complete_oauth_login(self, user_id: int, auth_code: str, timeout: int = 40) -> Tuple[bool, str]:
+        """
+        Completes the OAuth login flow by feeding the authorization code into the waiting AGY CLI process.
+        """
+        if user_id not in self.active_logins:
+            return False, "هیچ فرآیند لاگین فعالی برای شما یافت نشد. لطفاً دستور /login را مجدداً ارسال کنید."
+
+        login_data = self.active_logins[user_id]
+        proc = login_data["proc"]
+        backup_token = login_data.get("backup_token")
+        token_file = self.get_token_path(user_id)
+
+        clean_code = auth_code.strip()
+        clean_code = re.sub(r"^[`'\"]+|[`'\"]+$", "", clean_code)
+
+        try:
+            input_bytes = (clean_code + "\n").encode("utf-8")
+            proc.stdin.write(input_bytes)
+            await proc.stdin.drain()
+
+            stdout_data, stderr_data = await asyncio.wait_for(proc.communicate(), timeout=timeout)
+            out_str = stdout_data.decode("utf-8", errors="replace")
+            err_str = stderr_data.decode("utf-8", errors="replace")
+            logger.info(f"OAuth complete output for {user_id}: out='{out_str[:150]}' err='{err_str[:150]}'")
+
+            if token_file.exists() and token_file.stat().st_size > 10:
+                token_file.chmod(0o600)
+                self.active_logins.pop(user_id, None)
+                return True, "🎉 تبریک! ورود با حساب گوگل اختصاصی شما با موفقیت انجام و ذخیره شد."
+
+            if backup_token:
+                token_file.write_text(backup_token, encoding="utf-8")
+                token_file.chmod(0o600)
+
+            self.active_logins.pop(user_id, None)
+            return False, f"کد تایید معتبر نبود یا فرآیند با خطا مواجه شد:\n<pre>{escape_pre(err_str or out_str)}</pre>"
+
+        except asyncio.TimeoutError:
+            try:
+                proc.kill()
+            except Exception:
+                pass
+            self.active_logins.pop(user_id, None)
+            if backup_token:
+                token_file.write_text(backup_token, encoding="utf-8")
+                token_file.chmod(0o600)
+            return False, "زمان تایید کد احراز هویت به پایان رسید (Timeout). لطفاً مجدداً با /login تلاش کنید."
+
+        except Exception as e:
+            logger.error(f"Error completing OAuth login for user {user_id}: {e}", exc_info=True)
+            self.active_logins.pop(user_id, None)
+            return False, f"خطا در ثبت کد احراز هویت: {e}"
+
+    def cancel_oauth_login(self, user_id: int) -> bool:
+        """Cancels any pending login flow for a user."""
+        login_data = self.active_logins.pop(user_id, None)
+        if login_data:
+            proc = login_data.get("proc")
+            if proc:
+                try:
+                    proc.kill()
+                except Exception:
+                    pass
+            backup_token = login_data.get("backup_token")
+            token_file = self.get_token_path(user_id)
+            if backup_token and not token_file.exists():
+                try:
+                    token_file.write_text(backup_token, encoding="utf-8")
+                    token_file.chmod(0o600)
+                except Exception:
+                    pass
+            return True
+        return False
+
+    def import_token(self, user_id: int, token_data_str: str) -> Tuple[bool, str]:
+        """
+        Directly imports an existing antigravity-oauth-token content for the user.
+        """
+        raw_text = token_data_str.strip()
+        try:
+            parsed = json.loads(raw_text)
+            if not isinstance(parsed, dict) or ("token" not in parsed and "access_token" not in parsed and "refresh_token" not in parsed):
+                return False, "فرمت فایل یا متن توکن معتبر نیست. توکن باید فایل JSON حاوی اطلاعات احراز هویت AGY باشد."
+        except Exception:
+            return False, "متن ارسالی یک ساختار JSON معتبر برای توکن AGY نیست."
+
+        config_dir = self.get_config_dir(user_id)
+        config_dir.mkdir(parents=True, exist_ok=True)
+        token_file = self.get_token_path(user_id)
+
+        try:
+            token_file.write_text(raw_text, encoding="utf-8")
+            token_file.chmod(0o600)
+            return True, "✅ فایل توکن حساب اختصاصی شما با موفقیت ثبت و فعال شد."
+        except Exception as e:
+            return False, f"خطا در ذخیره‌سازی توکن: {e}"
+
+    def logout_user(self, user_id: int) -> Tuple[bool, str]:
+        """
+        Logs out a user by removing their custom token, falling back to server default.
+        """
+        self.cancel_oauth_login(user_id)
+        token_file = self.get_token_path(user_id)
+        if token_file.exists():
+            try:
+                token_file.unlink()
+                return True, "حساب اختصاصی شما با موفقیت خارج شد. از این پس درخواست‌های شما با حساب پیش‌فرض سرور اجرا می‌شوند."
+            except Exception as e:
+                return False, f"خطا در حذف توکن اختصاصی: {e}"
+        return True, "شما از قبل در حال استفاده از حساب پیش‌فرض سرور هستید."
+
+    def copy_master_token_to_user(self, user_id: int) -> Tuple[bool, str]:
+        """
+        Copies the current active server master token to a specific user's profile.
+        """
+        if not MASTER_TOKEN_PATH.exists():
+            return False, "توکن فعال سرور یافت نشد."
+
+        config_dir = self.get_config_dir(user_id)
+        config_dir.mkdir(parents=True, exist_ok=True)
+        dst = self.get_token_path(user_id)
+
+        try:
+            shutil.copy2(MASTER_TOKEN_PATH, dst)
+            dst.chmod(0o600)
+            return True, "توکن فعال سرور با موفقیت در پروفایل شما ثبت گردید."
+        except Exception as e:
+            return False, f"خطا در کپی توکن سرور: {e}"
+
+auth_manager = AuthManager()

telegram-agy-bot/bot.py

@@ -75,6 +75,7 @@ from remote_audit import remote_audit
 from sys_monitor import render_server_hardware_report
 from usage_monitor import fetch_and_render_usage_report
 from bot_actions import process_all_ai_actions
+from auth_manager import auth_manager
 from web_uploader import create_upload_token, get_upload_url, start_web_uploader_server, set_telegram_app
 
 LANG_ALIASES = {
@@ -217,11 +218,16 @@ def build_main_dashboard(chat_id: int) -> tuple[str, InlineKeyboardMarkup]:
     ctx_info = f"• 📏 <b>طول کانتکست:</b> <code>{curr_proj.last_context_length:,} توکن</code>\n" if curr_proj.last_context_length else ""
     ctx_info_en = f"• 📏 <b>Context Length:</b> <code>{curr_proj.last_context_length:,} tokens</code>\n" if curr_proj.last_context_length else ""
 
+    acc_status = auth_manager.get_account_status(chat_id)
+    acc_badge_fa = "🟢 اختصاصی" if acc_status["has_custom_account"] else "🟡 پیش‌فرض سرور"
+    acc_badge_en = "🟢 Custom" if acc_status["has_custom_account"] else "🟡 Server Default"
+
     if is_fa:
         text = (
             f"🎛 <b>کنترل‌پنل مدیریت Antigravity (AGY)</b>\n\n"
             f"• 📁 <b>پروژه فعال:</b> <code>{escape_html(curr_proj.name)}</code> {owner_str}\n"
             f"• 📂 <b>مسیر کاری:</b> <code>{escape_html(curr_proj.workspace)}</code>\n"
+            f"• 🔑 <b>حساب AGY:</b> <b>{acc_badge_fa}</b>\n"
             f"• 🧠 <b>مدل هوش مصنوعی:</b> <code>{curr_proj.model}</code>\n"
             f"• ⚡ <b>سطح استدلال:</b> <code>{curr_proj.effort}</code>\n"
             f"• 🌐 <b>زبان پاسخ‌دهی:</b> <code>{get_lang_display(curr_proj.language)}</code>\n"
@@ -238,7 +244,8 @@ def build_main_dashboard(chat_id: int) -> tuple[str, InlineKeyboardMarkup]:
             ],
             [
                 InlineKeyboardButton("🧠 حافظه هوش مصنوعی", callback_data="btn_memory_menu"),
-                InlineKeyboardButton("📈 سهمیه مصرف (AGY)", callback_data="btn_usage_menu"),
+                InlineKeyboardButton("📈 سهمیه مصرف", callback_data="btn_usage_menu"),
+                InlineKeyboardButton("🔑 حساب AGY", callback_data="btn_account_menu"),
             ],
             [
                 InlineKeyboardButton("🖥 سخت‌افزار سرور", callback_data="btn_hw_menu"),
@@ -261,6 +268,7 @@ def build_main_dashboard(chat_id: int) -> tuple[str, InlineKeyboardMarkup]:
             f"🎛 <b>Antigravity (AGY) Control Panel</b>\n\n"
             f"• 📁 <b>Active Project:</b> <code>{escape_html(curr_proj.name)}</code> {owner_str_en}\n"
             f"• 📂 <b>Workspace:</b> <code>{escape_html(curr_proj.workspace)}</code>\n"
+            f"• 🔑 <b>AGY Account:</b> <b>{acc_badge_en}</b>\n"
             f"• 🧠 <b>AI Model:</b> <code>{curr_proj.model}</code>\n"
             f"• ⚡ <b>Reasoning Effort:</b> <code>{curr_proj.effort}</code>\n"
             f"• 🌐 <b>Language:</b> <code>{get_lang_display(curr_proj.language)}</code>\n"
@@ -278,6 +286,7 @@ def build_main_dashboard(chat_id: int) -> tuple[str, InlineKeyboardMarkup]:
             [
                 InlineKeyboardButton("🧠 AI Memory", callback_data="btn_memory_menu"),
                 InlineKeyboardButton("📈 Quota & Usage", callback_data="btn_usage_menu"),
+                InlineKeyboardButton("🔑 AGY Account", callback_data="btn_account_menu"),
             ],
             [
                 InlineKeyboardButton("🖥 Server Hardware", callback_data="btn_hw_menu"),
@@ -975,6 +984,9 @@ def build_settings_menu(chat_id: int) -> tuple[str, InlineKeyboardMarkup]:
                 InlineKeyboardButton("🌐 زبان پاسخ‌دهی", callback_data="btn_lang_menu"),
                 InlineKeyboardButton("📊 وضعیت و کانتکست", callback_data="btn_status_menu"),
             ],
+            [
+                InlineKeyboardButton("🔑 حساب کاربری AGY (لاگین گوگل)", callback_data="btn_account_menu"),
+            ],
             [
                 InlineKeyboardButton("📦 دریافت فایل بکاپ (Zip)", callback_data=f"proj_backup:{curr_proj.name}"),
                 InlineKeyboardButton("🤝 اشتراک‌گذاری پروژه", callback_data="proj_share_menu"),
@@ -1007,6 +1019,9 @@ def build_settings_menu(chat_id: int) -> tuple[str, InlineKeyboardMarkup]:
                 InlineKeyboardButton("🌐 Response Language", callback_data="btn_lang_menu"),
                 InlineKeyboardButton("📊 Status & Tokens", callback_data="btn_status_menu"),
             ],
+            [
+                InlineKeyboardButton("🔑 AGY Account & Login", callback_data="btn_account_menu"),
+            ],
             [
                 InlineKeyboardButton("📦 Project Backup (Zip)", callback_data=f"proj_backup:{curr_proj.name}"),
                 InlineKeyboardButton("🤝 Share Project", callback_data="proj_share_menu"),
@@ -1018,24 +1033,124 @@ def build_settings_menu(chat_id: int) -> tuple[str, InlineKeyboardMarkup]:
 
     return text, InlineKeyboardMarkup(keyboard)
 
+def build_account_menu(chat_id: int) -> tuple[str, InlineKeyboardMarkup]:
+    """Generates the interactive AGY Account Management submenu."""
+    session = session_manager.get_or_create(chat_id)
+    is_fa = (session.language or "").lower() in ("fa", "farsi", "persian", "🇮🇷 persian / farsi (فارسی)")
+    is_admin_user = settings.is_admin(chat_id)
+    status = auth_manager.get_account_status(chat_id)
+
+    has_custom = status["has_custom_account"]
+    login_in_progress = status["is_login_in_progress"]
+
+    if is_fa:
+        if has_custom:
+            acc_type_badge = "🟢 <b>حساب اختصاصی فعال (Custom Account)</b>"
+            desc = (
+                "درخواست‌ها و مصرف مدل‌های هوش مصنوعی شما مستقیماً از <b>حساب گوگل اختصاصی خودتان</b> کسر می‌شود.\n"
+                f"• 📅 <i>تاریخ ثبت/بروزرسانی توکن:</i> <code>{status['last_modified']}</code>\n"
+            )
+        else:
+            acc_type_badge = "🟡 <b>حساب پیش‌فرض سرور (Server Shared)</b>"
+            desc = (
+                "شما هم‌اکنون از حساب عمومی سرور استفاده می‌کنید.\n"
+                "برای استفاده از سهمیه‌های اختصاصی، اعتبارات شخصی و حفظ کامل حریم کوتای خود، می‌توانید با حساب گوگل خود لاگین کنید.\n"
+            )
+
+        text = (
+            f"🔑 <b>مدیریت حساب کاربری هوش مصنوعی (AGY Account)</b>\n\n"
+            f"• 🆔 <b>شناسه کاربر:</b> <code>{chat_id}</code>\n"
+            f"• 👤 <b>وضعیت حساب:</b> {acc_type_badge}\n\n"
+            f"{desc}\n"
+            f"💡 <i>روش‌های ورود و اتصال:</i>\n"
+            f"۱. <b>ورود اینتراکتیو با گوگل:</b> دکمه «🔑 ورود به حساب گوگل» را بزنید.\n"
+            f"۲. <b>ارسال مستقیم فایل یا متن توکن:</b> می‌توانید فایل <code>antigravity-oauth-token</code> یا محتوای JSON توکن خود را در چت ارسال کنید."
+        )
+
+        buttons = []
+        if login_in_progress:
+            buttons.append([
+                InlineKeyboardButton("❌ لغو فرآیند لاگین جاری", callback_data="btn_account_cancel_login"),
+            ])
+        else:
+            buttons.append([
+                InlineKeyboardButton("🔑 ورود با حساب گوگل (OAuth)", callback_data="btn_account_start_login"),
+            ])
+
+        if has_custom:
+            buttons.append([
+                InlineKeyboardButton("🚪 خروج از حساب شخصی", callback_data="btn_account_logout"),
+            ])
+        elif is_admin_user:
+            buttons.append([
+                InlineKeyboardButton("📋 کپی توکن فعال سرور برای من", callback_data="btn_account_copy_server"),
+            ])
+
+        buttons.append([
+            InlineKeyboardButton("📊 استعلام سهمیه حساب", callback_data="btn_usage_menu"),
+            InlineKeyboardButton("⚙️ تنظیمات", callback_data="btn_settings_menu"),
+        ])
+        buttons.append([
+            InlineKeyboardButton("🏠 منوی اصلی", callback_data="btn_dashboard"),
+        ])
+    else:
+        if has_custom:
+            acc_type_badge = "🟢 <b>Active Custom Account</b>"
+            desc = f"All AI interactions are using your personal Google account quota.\n• 📅 <i>Last modified:</i> <code>{status['last_modified']}</code>\n"
+        else:
+            acc_type_badge = "🟡 <b>Server Default Account (Shared)</b>"
+            desc = "You are using the shared server account.\nLog in with your own Google account for isolated personal quotas.\n"
+
+        text = (
+            f"🔑 <b>AGY Account Management</b>\n\n"
+            f"• 🆔 <b>User ID:</b> <code>{chat_id}</code>\n"
+            f"• 👤 <b>Account Status:</b> {acc_type_badge}\n\n"
+            f"{desc}\n"
+            f"💡 <i>Login options:</i>\n"
+            f"1. <b>Google OAuth:</b> Tap 'Log in with Google' button below.\n"
+            f"2. <b>Direct Token:</b> Upload your <code>antigravity-oauth-token</code> file or JSON text."
+        )
+
+        buttons = []
+        if login_in_progress:
+            buttons.append([InlineKeyboardButton("❌ Cancel Pending Login", callback_data="btn_account_cancel_login")])
+        else:
+            buttons.append([InlineKeyboardButton("🔑 Log in with Google", callback_data="btn_account_start_login")])
+
+        if has_custom:
+            buttons.append([InlineKeyboardButton("🚪 Log out from Custom Account", callback_data="btn_account_logout")])
+        elif is_admin_user:
+            buttons.append([InlineKeyboardButton("📋 Copy Server Master Token to Me", callback_data="btn_account_copy_server")])
+
+        buttons.append([
+            InlineKeyboardButton("📊 Check Quota", callback_data="btn_usage_menu"),
+            InlineKeyboardButton("⚙️ Settings", callback_data="btn_settings_menu"),
+        ])
+        buttons.append([
+            InlineKeyboardButton("🏠 Dashboard", callback_data="btn_dashboard"),
+        ])
+
+    return text, InlineKeyboardMarkup(buttons)
+
 async def build_usage_report(chat_id: int) -> tuple[str, InlineKeyboardMarkup]:
     """Queries AGY CLI /usage and /credits and generates an interactive report with visual progress bars."""
     session = session_manager.get_or_create(chat_id)
     is_fa = (session.language or "").lower() in ("fa", "farsi", "persian", "🇮🇷 persian / farsi (فارسی)")
 
-    text = await fetch_and_render_usage_report(is_fa=is_fa)
+    text = await fetch_and_render_usage_report(is_fa=is_fa, user_id=chat_id)
 
     if is_fa:
         keyboard = [
             [
                 InlineKeyboardButton("🔄 بروزرسانی مصرف", callback_data="btn_usage_refresh"),
+                InlineKeyboardButton("🔑 مدیریت حساب AGY", callback_data="btn_account_menu"),
+            ],
+            [
                 InlineKeyboardButton("🖥 سخت‌افزار سرور", callback_data="btn_hw_menu"),
+                InlineKeyboardButton("📁 پروژه‌ها", callback_data="proj_menu"),
             ],
             [
-                InlineKeyboardButton("📁 مدیریت پروژه‌ها", callback_data="proj_menu"),
                 InlineKeyboardButton("🏠 منوی اصلی", callback_data="btn_dashboard"),
-            ],
-            [
                 InlineKeyboardButton("🔙 بستن", callback_data="proj_close"),
             ],
         ]
@@ -1043,13 +1158,14 @@ async def build_usage_report(chat_id: int) -> tuple[str, InlineKeyboardMarkup]:
         keyboard = [
             [
                 InlineKeyboardButton("🔄 Refresh Quota", callback_data="btn_usage_refresh"),
+                InlineKeyboardButton("🔑 AGY Account", callback_data="btn_account_menu"),
+            ],
+            [
                 InlineKeyboardButton("🖥 Server Hardware", callback_data="btn_hw_menu"),
-            ],
-            [
                 InlineKeyboardButton("📁 Projects", callback_data="proj_menu"),
-                InlineKeyboardButton("🏠 Main Dashboard", callback_data="btn_dashboard"),
             ],
             [
+                InlineKeyboardButton("🏠 Main Dashboard", callback_data="btn_dashboard"),
                 InlineKeyboardButton("🔙 Close", callback_data="proj_close"),
             ],
         ]
@@ -4038,6 +4154,74 @@ async def usage_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
     text, markup = await build_usage_report(chat_id)
     await status_msg.edit_text(text, parse_mode=constants.ParseMode.HTML, reply_markup=markup)
 
+# Command: /account or /profile_auth or /acc
+@check_auth
+async def account_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    chat_id = update.effective_chat.id
+    text, markup = build_account_menu(chat_id)
+    await update.message.reply_html(text, reply_markup=markup)
+
+# Command: /login or /signin
+@check_auth
+async def login_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    chat_id = update.effective_chat.id
+    session = session_manager.get_or_create(chat_id)
+    is_fa = (session.language or "").lower() in ("fa", "farsi", "persian", "🇮🇷 persian / farsi (فارسی)")
+
+    if context.args:
+        # User provided auth code directly: /login 4/0A...
+        auth_code = " ".join(context.args).strip()
+        wait_msg = await update.message.reply_html("⏳ <i>در حال ثبت و بررسی کد احراز هویت گوگل...</i>" if is_fa else "⏳ <i>Submitting and verifying Google Auth code...</i>")
+        success, msg = await auth_manager.complete_oauth_login(chat_id, auth_code)
+        if success:
+            text, markup = build_account_menu(chat_id)
+            await wait_msg.edit_text(f"{msg}\n\n{text}", parse_mode=constants.ParseMode.HTML, reply_markup=markup)
+        else:
+            await wait_msg.edit_text(msg, parse_mode=constants.ParseMode.HTML)
+        return
+
+    wait_msg = await update.message.reply_html("⏳ <i>در حال ایجاد نشست و لینک ورود اختصاصی گوگل...</i>" if is_fa else "⏳ <i>Generating Google OAuth login URL...</i>")
+    success, msg, auth_url = await auth_manager.start_oauth_login(chat_id)
+
+    if not success or not auth_url:
+        await wait_msg.edit_text(f"❌ {msg}", parse_mode=constants.ParseMode.HTML)
+        return
+
+    if is_fa:
+        prompt_text = (
+            "🔑 <b>ورود به حساب اختصاصی گوگل (AGY OAuth Login)</b>\n\n"
+            "۱. روی دکمه زیر کلیک کنید تا صفحه رسمی ورود به حساب گوگل باز شود.\n"
+            "۲. وارد حساب کاربری گوگل خود شوید و اجازه دسترسی را تایید نمایید.\n"
+            "۳. <b>کد تایید احراز هویت (Authorization Code)</b> نمایش داده شده را کپی کرده و در پاسخ به همین پیام برای ربات ارسال کنید (یا دستور <code>/login &lt;کد&gt;</code> را بفرستید).\n\n"
+            "⏳ <i>فرصت ارسال کد: ۶۰ ثانیه</i>"
+        )
+        keyboard = [
+            [InlineKeyboardButton("🌐 ورود به حساب کاربری گوگل", url=auth_url)],
+            [InlineKeyboardButton("❌ انصراف / لغو", callback_data="btn_account_cancel_login")],
+        ]
+    else:
+        prompt_text = (
+            "🔑 <b>Google Account Login (AGY OAuth)</b>\n\n"
+            "1. Click the button below to open Google's authorization page.\n"
+            "2. Sign in to your Google Account and grant access permissions.\n"
+            "3. Copy the <b>Authorization Code</b> and paste it as a message to this bot.\n\n"
+            "⏳ <i>Timeout: 60 seconds</i>"
+        )
+        keyboard = [
+            [InlineKeyboardButton("🌐 Sign in with Google", url=auth_url)],
+            [InlineKeyboardButton("❌ Cancel", callback_data="btn_account_cancel_login")],
+        ]
+
+    await wait_msg.edit_text(prompt_text, parse_mode=constants.ParseMode.HTML, reply_markup=InlineKeyboardMarkup(keyboard), disable_web_page_preview=True)
+
+# Command: /logout or /signout
+@check_auth
+async def logout_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    chat_id = update.effective_chat.id
+    success, msg = auth_manager.logout_user(chat_id)
+    text, markup = build_account_menu(chat_id)
+    await update.message.reply_html(f"{'✅' if success else '❌'} {msg}\n\n{text}", reply_markup=markup)
+
 # Command: /memory or /memories
 @check_auth
 async def memory_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
@@ -5732,6 +5916,76 @@ async def callback_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
         except Exception:
             await query.message.reply_html(text, reply_markup=markup)
 
+    elif data == "btn_account_menu":
+        text, markup = build_account_menu(chat_id)
+        try:
+            await query.edit_message_text(text, parse_mode=constants.ParseMode.HTML, reply_markup=markup)
+        except Exception:
+            await query.message.reply_html(text, reply_markup=markup)
+
+    elif data == "btn_account_start_login":
+        await query.answer("⏳ در حال ایجاد نشست احراز هویت گوگل..." if is_fa else "⏳ Starting Google OAuth session...")
+        success, msg, auth_url = await auth_manager.start_oauth_login(chat_id)
+        if not success or not auth_url:
+            await query.message.reply_html(f"❌ {msg}")
+            return
+
+        if is_fa:
+            prompt_text = (
+                "🔑 <b>ورود به حساب اختصاصی گوگل (AGY OAuth Login)</b>\n\n"
+                "۱. روی دکمه زیر کلیک کنید تا صفحه رسمی گوگل باز شود.\n"
+                "۲. وارد حساب کاربری گوگل خود شوید و اجازه دسترسی را تایید نمایید.\n"
+                "۳. <b>کد تایید احراز هویت (Authorization Code)</b> نمایش داده شده را کپی کرده و در چت ارسال کنید.\n\n"
+                "⏳ <i>فرصت ارسال کد: ۶۰ ثانیه</i>"
+            )
+            keyboard = [
+                [InlineKeyboardButton("🌐 ورود به حساب کاربری گوگل", url=auth_url)],
+                [InlineKeyboardButton("❌ انصراف / لغو", callback_data="btn_account_cancel_login")],
+            ]
+        else:
+            prompt_text = (
+                "🔑 <b>Google Account Login (AGY OAuth)</b>\n\n"
+                "1. Click the button below to open Google's authorization page.\n"
+                "2. Sign in to your Google Account and grant access permissions.\n"
+                "3. Copy the <b>Authorization Code</b> and paste it in chat.\n\n"
+                "⏳ <i>Timeout: 60 seconds</i>"
+            )
+            keyboard = [
+                [InlineKeyboardButton("🌐 Sign in with Google", url=auth_url)],
+                [InlineKeyboardButton("❌ Cancel", callback_data="btn_account_cancel_login")],
+            ]
+        await query.message.reply_html(prompt_text, reply_markup=InlineKeyboardMarkup(keyboard), disable_web_page_preview=True)
+
+    elif data == "btn_account_cancel_login":
+        auth_manager.cancel_oauth_login(chat_id)
+        await query.answer("فرآیند لاگین لغو شد." if is_fa else "Login canceled.")
+        text, markup = build_account_menu(chat_id)
+        try:
+            await query.edit_message_text(text, parse_mode=constants.ParseMode.HTML, reply_markup=markup)
+        except Exception:
+            await query.message.reply_html(text, reply_markup=markup)
+
+    elif data == "btn_account_logout":
+        success, msg = auth_manager.logout_user(chat_id)
+        await query.answer(msg, show_alert=True)
+        text, markup = build_account_menu(chat_id)
+        try:
+            await query.edit_message_text(text, parse_mode=constants.ParseMode.HTML, reply_markup=markup)
+        except Exception:
+            await query.message.reply_html(text, reply_markup=markup)
+
+    elif data == "btn_account_copy_server":
+        if not is_admin_user:
+            await query.answer("⛔ فقط مدیر مجاز به کپی توکن سرور است.", show_alert=True)
+            return
+        success, msg = auth_manager.copy_master_token_to_user(chat_id)
+        await query.answer(msg, show_alert=True)
+        text, markup = build_account_menu(chat_id)
+        try:
+            await query.edit_message_text(text, parse_mode=constants.ParseMode.HTML, reply_markup=markup)
+        except Exception:
+            await query.message.reply_html(text, reply_markup=markup)
+
     elif data == "btn_status_menu":
         accessible = session_manager.get_all_accessible_projects(chat_id)
         if not curr_proj:
@@ -7095,6 +7349,33 @@ async def message_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
     session = session_manager.get_or_create(chat_id)
     is_fa = (session.language or "").lower() in ("fa", "farsi", "persian", "🇮🇷 persian / farsi (فارسی)")
 
+    # 0. Check if user has an active OAuth login waiting for authorization code
+    if chat_id in auth_manager.active_logins:
+        code_input = update.message.text.strip()
+        if code_input.lower() in ("/cancel", "لغو", "انصراف", "cancel"):
+            auth_manager.cancel_oauth_login(chat_id)
+            text, markup = build_account_menu(chat_id)
+            await update.message.reply_html("❌ فرآیند لاگین گوگل لغو گردید.\n\n" + text, reply_markup=markup)
+            return
+
+        wait_msg = await update.message.reply_html("⏳ <i>در حال ثبت و فعال‌سازی حساب کاربری با کد دریافتی...</i>" if is_fa else "⏳ <i>Completing login with authorization code...</i>")
+        success, msg = await auth_manager.complete_oauth_login(chat_id, code_input)
+        if success:
+            text, markup = build_account_menu(chat_id)
+            await wait_msg.edit_text(f"{msg}\n\n{text}", parse_mode=constants.ParseMode.HTML, reply_markup=markup)
+        else:
+            await wait_msg.edit_text(msg, parse_mode=constants.ParseMode.HTML)
+        return
+
+    # Check if user sent a JSON token directly in text
+    raw_stripped = update.message.text.strip()
+    if raw_stripped.startswith("{") and raw_stripped.endswith("}") and ("token" in raw_stripped or "access_token" in raw_stripped):
+        success, msg = auth_manager.import_token(chat_id, raw_stripped)
+        if success:
+            text, markup = build_account_menu(chat_id)
+            await update.message.reply_html(f"{msg}\n\n{text}", reply_markup=markup)
+            return
+
     # 1. Check if user is in an interactive input state (e.g. configuring FTP fields)
     if chat_id in USER_INPUT_STATES:
         state = USER_INPUT_STATES.pop(chat_id)
@@ -7287,6 +7568,22 @@ async def file_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
     session = session_manager.get_or_create(chat_id)
     is_fa = (session.language or "").lower() in ("fa", "farsi", "persian", "🇮🇷 persian / farsi (فارسی)")
 
+    # Check if uploaded document is antigravity-oauth-token or a token json
+    if message.document:
+        doc_name = (message.document.file_name or "").lower()
+        if "oauth-token" in doc_name or doc_name.endswith(".token") or (doc_name.endswith(".json") and "token" in doc_name):
+            try:
+                doc_file = await message.document.get_file()
+                content_bytes = await doc_file.download_as_bytearray()
+                content_str = content_bytes.decode("utf-8", errors="ignore")
+                success, msg = auth_manager.import_token(chat_id, content_str)
+                if success:
+                    text, markup = build_account_menu(chat_id)
+                    await update.message.reply_html(f"{msg}\n\n{text}", reply_markup=markup)
+                    return
+            except Exception as ex:
+                logger.error(f"Failed to process uploaded token file: {ex}")
+
     if not curr_proj:
         msg = (
             "⚠️ <b>شما هنوز هیچ پروژه‌ای ایجاد نکرده‌اید!</b>\n\n"
@@ -7816,6 +8113,9 @@ async def on_startup(app: Application):
             BotCommand("new", "گفتگوی جدید / New chat"),
             BotCommand("compact", "فشرده‌سازی کانتکست / Compact context"),
             BotCommand("usage", "سهمیه و مصرف / Quota & usage"),
+            BotCommand("account", "حساب کاربری AGY / AGY Account"),
+            BotCommand("login", "ورود به حساب گوگل / Google Login"),
+            BotCommand("logout", "خروج از حساب / Logout"),
             BotCommand("model", "تغییر مدل / AI model"),
             BotCommand("lang", "تنظیم زبان / Set language"),
             BotCommand("effort", "سطح استدلال / Reasoning effort"),
@@ -8012,6 +8312,9 @@ def main():
     application.add_handler(CommandHandler(["memory", "memories", "hafeze"], memory_command))
     application.add_handler(CommandHandler("help", help_command))
     application.add_handler(CommandHandler(["usage", "quota", "credits"], usage_command))
+    application.add_handler(CommandHandler(["account", "profile_auth", "acc"], account_command))
+    application.add_handler(CommandHandler(["login", "google_login", "signin"], login_command))
+    application.add_handler(CommandHandler(["logout", "signout"], logout_command))
     application.add_handler(CommandHandler(["new", "reset"], new_command))
     application.add_handler(CommandHandler(["compact", "compress", "fashorde"], compact_command))
     application.add_handler(CommandHandler(["lastconv", "lastconversation", "openlast", "resume", "open_last"], last_conversation_command))

telegram-agy-bot/usage_monitor.py

@@ -1,3 +1,4 @@
+import os
 import re
 import time
 import asyncio
@@ -122,8 +123,10 @@ def parse_agy_usage_output(usage_raw: str, credits_raw: str) -> Dict[str, Any]:
 
     return data
 
-async def fetch_and_render_usage_report(is_fa: bool = True) -> str:
+async def fetch_and_render_usage_report(is_fa: bool = True, user_id: Optional[int] = None) -> str:
     """Executes agy /usage and /credits and formats with visual progress bars."""
+    from auth_manager import auth_manager
+    env = auth_manager.get_user_env(user_id) if user_id else os.environ.copy()
     usage_raw = ""
     credits_raw = ""
     try:
@@ -131,6 +134,7 @@ async def fetch_and_render_usage_report(is_fa: bool = True) -> str:
             "agy", "--print", "/usage",
             stdout=asyncio.subprocess.PIPE,
             stderr=asyncio.subprocess.PIPE,
+            env=env,
         )
         out1, _ = await proc1.communicate()
         usage_raw = out1.decode("utf-8", errors="replace").strip()
@@ -139,6 +143,7 @@ async def fetch_and_render_usage_report(is_fa: bool = True) -> str:
             "agy", "--print", "/credits",
             stdout=asyncio.subprocess.PIPE,
             stderr=asyncio.subprocess.PIPE,
+            env=env,
         )
         out2, _ = await proc2.communicate()
         credits_raw = out2.decode("utf-8", errors="replace").strip()