diff --git a/bot/api/discord_auth.py b/bot/api/discord_auth.py index 1c207dd..e934595 100644 --- a/bot/api/discord_auth.py +++ b/bot/api/discord_auth.py @@ -13,7 +13,9 @@ from __future__ import annotations +import asyncio import re +import time from typing import Any import aiohttp @@ -26,6 +28,17 @@ ADMINISTRATOR = 0x8 GUILD_PATH_RE = re.compile(r"^/api/v1/guilds/(\d+)(?:/|$)") +# Short-lived process caches — cut Discord 429s from dashboard reloads +_USER_CACHE_TTL = 60.0 +_GUILDS_CACHE_TTL = 60.0 +_user_cache: dict[str, tuple[float, dict[str, Any]]] = {} +_guilds_cache: dict[str, tuple[float, list[dict[str, Any]]]] = {} +_guilds_inflight: dict[str, asyncio.Future[list[dict[str, Any]]]] = {} + + +def _token_key(access_token: str) -> str: + return f"{access_token[:12]}:{len(access_token)}:{access_token[-8:]}" + def _parse_admin_ids() -> list[str]: import os @@ -46,33 +59,92 @@ def can_manage_guild(guild: dict[str, Any]) -> bool: return bool(perms & ADMINISTRATOR) or bool(perms & MANAGE_GUILD) +def _discord_auth_error(status: int) -> HTTPException: + if status == 429: + return HTTPException( + status_code=429, + detail="Discord rate limit — please wait a few seconds and retry.", + ) + if status in (401, 403): + return HTTPException( + status_code=401, + detail="Invalid or expired Discord access token. Please sign out and sign in again.", + ) + return HTTPException( + status_code=403, + detail=f"Discord API error ({status}).", + ) + + async def fetch_user_guilds(access_token: str) -> list[dict[str, Any]]: + key = _token_key(access_token) + cached = _guilds_cache.get(key) + if cached and cached[0] > time.monotonic(): + return cached[1] + + existing = _guilds_inflight.get(key) + if existing is not None: + return await existing + + loop = asyncio.get_running_loop() + fut: asyncio.Future[list[dict[str, Any]]] = loop.create_future() + _guilds_inflight[key] = fut + + try: + guilds = await _fetch_user_guilds_uncached(access_token) + _guilds_cache[key] = (time.monotonic() + _GUILDS_CACHE_TTL, guilds) + fut.set_result(guilds) + return guilds + except Exception as exc: + fut.set_exception(exc) + raise + finally: + _guilds_inflight.pop(key, None) + + +async def _fetch_user_guilds_uncached(access_token: str) -> list[dict[str, Any]]: + max_attempts = 3 async with aiohttp.ClientSession() as session: - async with session.get( - "https://discord.com/api/users/@me/guilds", - headers={"Authorization": f"Bearer {access_token}"}, - ) as resp: - if resp.status != 200: - raise HTTPException( - status_code=403, - detail="Invalid or expired Discord access token.", - ) - data = await resp.json() - return data if isinstance(data, list) else [] + for attempt in range(1, max_attempts + 1): + async with session.get( + "https://discord.com/api/users/@me/guilds", + headers={"Authorization": f"Bearer {access_token}"}, + ) as resp: + if resp.status == 200: + data = await resp.json() + return data if isinstance(data, list) else [] + + if resp.status == 429 and attempt < max_attempts: + retry_after = 1.0 + try: + body = await resp.json() + retry_after = float(body.get("retry_after", 1)) + except Exception: + retry_after = float(resp.headers.get("Retry-After", "1") or 1) + await asyncio.sleep(min(max(retry_after, 0.3), 5.0)) + continue + + raise _discord_auth_error(resp.status) + + raise _discord_auth_error(429) async def verify_discord_token(access_token: str) -> dict[str, Any]: + key = _token_key(access_token) + cached = _user_cache.get(key) + if cached and cached[0] > time.monotonic(): + return cached[1] + async with aiohttp.ClientSession() as session: async with session.get( "https://discord.com/api/users/@me", headers={"Authorization": f"Bearer {access_token}"}, ) as resp: if resp.status != 200: - raise HTTPException( - status_code=403, - detail="Invalid or expired Discord access token.", - ) - return await resp.json() + raise _discord_auth_error(resp.status) + user = await resp.json() + _user_cache[key] = (time.monotonic() + _USER_CACHE_TTL, user) + return user async def user_can_manage_guild(access_token: str, guild_id: int) -> bool: diff --git a/bot/api/routes/guilds.py b/bot/api/routes/guilds.py index c3e8fb4..fe07c12 100644 --- a/bot/api/routes/guilds.py +++ b/bot/api/routes/guilds.py @@ -14,7 +14,7 @@ from utils.db_paths import db_path, jsondb_path from fastapi import APIRouter, Depends, HTTPException, Request from api.dependencies import get_bot, limiter -from api.discord_auth import require_discord_session, get_manageable_guild_ids +from api.discord_auth import get_manageable_guild_ids, get_discord_headers from api.db_manager import db_manager from api.schemas import ( GuildSummary, GuildDetails, PrefixConfig, AutomodConfig, @@ -46,8 +46,11 @@ router = APIRouter() async def list_guilds(request: Request, bot: "Axiom" = Depends(get_bot)): """ Lists guilds the bot is in, filtered to those the caller can manage on Discord. + Session is already verified by middleware — only fetch manageable guild IDs here. """ - token, _ = await require_discord_session(request) + token, _ = get_discord_headers(request) + if not token: + raise HTTPException(status_code=403, detail="X-Discord-Access-Token header is required.") manageable_ids = await get_manageable_guild_ids(token) guilds_list = [] diff --git a/bot/cogs/commands/ai.py b/bot/cogs/commands/ai.py index 50acbae..16bc96c 100644 --- a/bot/cogs/commands/ai.py +++ b/bot/cogs/commands/ai.py @@ -323,19 +323,19 @@ class AI (commands .Cog ): import os - db_path =db_path("ai_data.db") - if os .path .exists (db_path ): + ai_db = db_path("ai_data.db") + if os .path .exists (ai_db ): try : - test_conn =await aiosqlite .connect (db_path ) + test_conn =await aiosqlite .connect (ai_db ) await test_conn .execute ("SELECT name FROM sqlite_master WHERE type='table';") await test_conn .close () except Exception as e : - os .remove (db_path ) + os .remove (ai_db ) logger .info ("Removed corrupted AI database, creating new one") - self .bot .db =await aiosqlite .connect (db_path ) + self .bot .db =await aiosqlite .connect (ai_db ) logger .info ("AI database connection initialized") await self .bot .db .execute (""" @@ -401,12 +401,12 @@ class AI (commands .Cog ): import aiosqlite import os - db_path =db_path("ai_data.db") - if not os .path .exists (db_path ): + ai_db = db_path("ai_data.db") + if not os .path .exists (ai_db ): logger .info ("AI database doesn't exist, will be created on first use") return - self .bot .db =await aiosqlite .connect (db_path ) + self .bot .db =await aiosqlite .connect (ai_db ) logger .info ("AI database connection initialized for loading") diff --git a/dashboard/app/dashboard/guilds/page.tsx b/dashboard/app/dashboard/guilds/page.tsx index de8bcf1..11a5d41 100644 --- a/dashboard/app/dashboard/guilds/page.tsx +++ b/dashboard/app/dashboard/guilds/page.tsx @@ -51,46 +51,23 @@ export default async function GuildsPage() { redirect("/"); } - let botGuilds: GuildSummary[] = []; - let userGuilds: any[] = []; - let userDiscordError: string | null = null; + // Bot API already filters to guilds the user can manage — no second Discord @me/guilds call + let guilds: GuildSummary[] = []; let botError: string | null = null; try { - botGuilds = await api.listGuilds(); + guilds = await api.listGuilds(); } catch (err: any) { console.error("Failed to fetch bot guilds:", err); - botError = err.message || "Failed to load bot servers."; - } - - try { - const { fetchUserGuilds } = await import("@/lib/guild-auth"); - userGuilds = await fetchUserGuilds(session.accessToken); - } catch (err: any) { - console.error("Discord API Error:", err); - userDiscordError = err?.message || "Error connecting to Discord."; - } - - const MANAGE_GUILD = BigInt(0x20); - const ADMINISTRATOR = BigInt(0x8); - const adminUserGuilds = userGuilds.filter((g) => { - try { - const perms = BigInt(g.permissions); - return ( - (perms & ADMINISTRATOR) === ADMINISTRATOR || - (perms & MANAGE_GUILD) === MANAGE_GUILD || - g.owner === true - ); - } catch { - return g.owner === true; + const msg = typeof err?.message === "string" ? err.message : "Failed to load bot servers."; + botError = msg; + // Expired OAuth — force re-login + if (err?.status === 401 || /expired|sign in again/i.test(msg)) { + redirect("/"); } - }); - - const adminGuildIds = new Set(adminUserGuilds.map((g) => String(g.id))); - const guilds = botGuilds.filter((g) => adminGuildIds.has(String(g.id))); + } const inviteConfigured = Boolean(getBotInviteUrl()); - // Hard failure only when bot API is down — Discord sync issues still show invite CTA const hardError = botError; return ( @@ -132,11 +109,6 @@ export default async function GuildsPage() { Invite Axiom to a Discord server where you have Manage Server{" "} permissions, then come back here to configure it.
- {userDiscordError && ( -- Note: {userDiscordError} -
- )}