Enhance Discord authentication and guild management by implementing caching for user and guild data to reduce API calls and improve performance. Update error handling for Discord API responses and refactor the guilds listing logic in the dashboard to streamline data retrieval. Adjust retry logic for rate limits and improve session error handling in the authentication flow.
Some checks failed
CI / Bot (Python) (push) Successful in 12s
CI / Dashboard (Next.js) (push) Failing after 9s

This commit is contained in:
TheOnlyMace
2026-07-21 22:33:30 +02:00
parent 60390009d1
commit f15c869993
6 changed files with 131 additions and 73 deletions

View File

@@ -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:

View File

@@ -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 = []

View File

@@ -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")