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.
This commit is contained in:
@@ -13,7 +13,9 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import re
|
import re
|
||||||
|
import time
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import aiohttp
|
import aiohttp
|
||||||
@@ -26,6 +28,17 @@ ADMINISTRATOR = 0x8
|
|||||||
|
|
||||||
GUILD_PATH_RE = re.compile(r"^/api/v1/guilds/(\d+)(?:/|$)")
|
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]:
|
def _parse_admin_ids() -> list[str]:
|
||||||
import os
|
import os
|
||||||
@@ -46,33 +59,92 @@ def can_manage_guild(guild: dict[str, Any]) -> bool:
|
|||||||
return bool(perms & ADMINISTRATOR) or bool(perms & MANAGE_GUILD)
|
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]]:
|
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 aiohttp.ClientSession() as session:
|
||||||
async with session.get(
|
for attempt in range(1, max_attempts + 1):
|
||||||
"https://discord.com/api/users/@me/guilds",
|
async with session.get(
|
||||||
headers={"Authorization": f"Bearer {access_token}"},
|
"https://discord.com/api/users/@me/guilds",
|
||||||
) as resp:
|
headers={"Authorization": f"Bearer {access_token}"},
|
||||||
if resp.status != 200:
|
) as resp:
|
||||||
raise HTTPException(
|
if resp.status == 200:
|
||||||
status_code=403,
|
data = await resp.json()
|
||||||
detail="Invalid or expired Discord access token.",
|
return data if isinstance(data, list) else []
|
||||||
)
|
|
||||||
data = await resp.json()
|
if resp.status == 429 and attempt < max_attempts:
|
||||||
return data if isinstance(data, list) else []
|
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]:
|
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 aiohttp.ClientSession() as session:
|
||||||
async with session.get(
|
async with session.get(
|
||||||
"https://discord.com/api/users/@me",
|
"https://discord.com/api/users/@me",
|
||||||
headers={"Authorization": f"Bearer {access_token}"},
|
headers={"Authorization": f"Bearer {access_token}"},
|
||||||
) as resp:
|
) as resp:
|
||||||
if resp.status != 200:
|
if resp.status != 200:
|
||||||
raise HTTPException(
|
raise _discord_auth_error(resp.status)
|
||||||
status_code=403,
|
user = await resp.json()
|
||||||
detail="Invalid or expired Discord access token.",
|
_user_cache[key] = (time.monotonic() + _USER_CACHE_TTL, user)
|
||||||
)
|
return user
|
||||||
return await resp.json()
|
|
||||||
|
|
||||||
|
|
||||||
async def user_can_manage_guild(access_token: str, guild_id: int) -> bool:
|
async def user_can_manage_guild(access_token: str, guild_id: int) -> bool:
|
||||||
|
|||||||
@@ -14,7 +14,7 @@ from utils.db_paths import db_path, jsondb_path
|
|||||||
|
|
||||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||||
from api.dependencies import get_bot, limiter
|
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.db_manager import db_manager
|
||||||
from api.schemas import (
|
from api.schemas import (
|
||||||
GuildSummary, GuildDetails, PrefixConfig, AutomodConfig,
|
GuildSummary, GuildDetails, PrefixConfig, AutomodConfig,
|
||||||
@@ -46,8 +46,11 @@ router = APIRouter()
|
|||||||
async def list_guilds(request: Request, bot: "Axiom" = Depends(get_bot)):
|
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.
|
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)
|
manageable_ids = await get_manageable_guild_ids(token)
|
||||||
|
|
||||||
guilds_list = []
|
guilds_list = []
|
||||||
|
|||||||
@@ -323,19 +323,19 @@ class AI (commands .Cog ):
|
|||||||
import os
|
import os
|
||||||
|
|
||||||
|
|
||||||
db_path =db_path("ai_data.db")
|
ai_db = db_path("ai_data.db")
|
||||||
if os .path .exists (db_path ):
|
if os .path .exists (ai_db ):
|
||||||
try :
|
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 .execute ("SELECT name FROM sqlite_master WHERE type='table';")
|
||||||
await test_conn .close ()
|
await test_conn .close ()
|
||||||
except Exception as e :
|
except Exception as e :
|
||||||
|
|
||||||
os .remove (db_path )
|
os .remove (ai_db )
|
||||||
logger .info ("Removed corrupted AI database, creating new one")
|
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")
|
logger .info ("AI database connection initialized")
|
||||||
|
|
||||||
await self .bot .db .execute ("""
|
await self .bot .db .execute ("""
|
||||||
@@ -401,12 +401,12 @@ class AI (commands .Cog ):
|
|||||||
import aiosqlite
|
import aiosqlite
|
||||||
import os
|
import os
|
||||||
|
|
||||||
db_path =db_path("ai_data.db")
|
ai_db = db_path("ai_data.db")
|
||||||
if not os .path .exists (db_path ):
|
if not os .path .exists (ai_db ):
|
||||||
logger .info ("AI database doesn't exist, will be created on first use")
|
logger .info ("AI database doesn't exist, will be created on first use")
|
||||||
return
|
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")
|
logger .info ("AI database connection initialized for loading")
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -51,46 +51,23 @@ export default async function GuildsPage() {
|
|||||||
redirect("/");
|
redirect("/");
|
||||||
}
|
}
|
||||||
|
|
||||||
let botGuilds: GuildSummary[] = [];
|
// Bot API already filters to guilds the user can manage — no second Discord @me/guilds call
|
||||||
let userGuilds: any[] = [];
|
let guilds: GuildSummary[] = [];
|
||||||
let userDiscordError: string | null = null;
|
|
||||||
let botError: string | null = null;
|
let botError: string | null = null;
|
||||||
|
|
||||||
try {
|
try {
|
||||||
botGuilds = await api.listGuilds();
|
guilds = await api.listGuilds();
|
||||||
} catch (err: any) {
|
} catch (err: any) {
|
||||||
console.error("Failed to fetch bot guilds:", err);
|
console.error("Failed to fetch bot guilds:", err);
|
||||||
botError = err.message || "Failed to load bot servers.";
|
const msg = typeof err?.message === "string" ? err.message : "Failed to load bot servers.";
|
||||||
}
|
botError = msg;
|
||||||
|
// Expired OAuth — force re-login
|
||||||
try {
|
if (err?.status === 401 || /expired|sign in again/i.test(msg)) {
|
||||||
const { fetchUserGuilds } = await import("@/lib/guild-auth");
|
redirect("/");
|
||||||
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 adminGuildIds = new Set(adminUserGuilds.map((g) => String(g.id)));
|
|
||||||
const guilds = botGuilds.filter((g) => adminGuildIds.has(String(g.id)));
|
|
||||||
|
|
||||||
const inviteConfigured = Boolean(getBotInviteUrl());
|
const inviteConfigured = Boolean(getBotInviteUrl());
|
||||||
// Hard failure only when bot API is down — Discord sync issues still show invite CTA
|
|
||||||
const hardError = botError;
|
const hardError = botError;
|
||||||
|
|
||||||
return (
|
return (
|
||||||
@@ -132,11 +109,6 @@ export default async function GuildsPage() {
|
|||||||
Invite Axiom to a Discord server where you have <span className="text-slate-200">Manage Server</span>{" "}
|
Invite Axiom to a Discord server where you have <span className="text-slate-200">Manage Server</span>{" "}
|
||||||
permissions, then come back here to configure it.
|
permissions, then come back here to configure it.
|
||||||
</p>
|
</p>
|
||||||
{userDiscordError && (
|
|
||||||
<p className="text-amber-400/90 text-sm mt-4 max-w-md mx-auto">
|
|
||||||
Note: {userDiscordError}
|
|
||||||
</p>
|
|
||||||
)}
|
|
||||||
<div className="mt-8 flex flex-col sm:flex-row gap-3 justify-center items-center">
|
<div className="mt-8 flex flex-col sm:flex-row gap-3 justify-center items-center">
|
||||||
{inviteConfigured ? (
|
{inviteConfigured ? (
|
||||||
<AddBotButton className="px-8 h-12 text-base shadow-lg shadow-primary/25" />
|
<AddBotButton className="px-8 h-12 text-base shadow-lg shadow-primary/25" />
|
||||||
|
|||||||
@@ -79,8 +79,14 @@ export const authOptions: AuthOptions = {
|
|||||||
}
|
}
|
||||||
|
|
||||||
const expiresAt = typeof token.expiresAt === "number" ? token.expiresAt : 0;
|
const expiresAt = typeof token.expiresAt === "number" ? token.expiresAt : 0;
|
||||||
// Refresh ~60s before expiry
|
// Refresh when expired, missing expiry, or previous refresh failed
|
||||||
if (Date.now() < expiresAt * 1000 - 60_000) {
|
const stillValid =
|
||||||
|
Boolean(token.accessToken) &&
|
||||||
|
expiresAt > 0 &&
|
||||||
|
Date.now() < expiresAt * 1000 - 60_000 &&
|
||||||
|
!token.error;
|
||||||
|
|
||||||
|
if (stillValid) {
|
||||||
return token;
|
return token;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -100,6 +106,11 @@ export const authOptions: AuthOptions = {
|
|||||||
// @ts-ignore
|
// @ts-ignore
|
||||||
session.error = token.error;
|
session.error = token.error;
|
||||||
}
|
}
|
||||||
|
// Surface refresh failures so pages can redirect to sign-in
|
||||||
|
if (token.error) {
|
||||||
|
// @ts-ignore
|
||||||
|
session.error = token.error;
|
||||||
|
}
|
||||||
return session;
|
return session;
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -14,8 +14,8 @@ export interface DiscordGuild {
|
|||||||
const MANAGE_GUILD = BigInt(0x20);
|
const MANAGE_GUILD = BigInt(0x20);
|
||||||
const ADMINISTRATOR = BigInt(0x8);
|
const ADMINISTRATOR = BigInt(0x8);
|
||||||
|
|
||||||
const CACHE_TTL_MS = 60_000;
|
const CACHE_TTL_MS = 90_000;
|
||||||
const MAX_RETRIES = 3;
|
const MAX_RETRIES = 4;
|
||||||
|
|
||||||
type CacheEntry = { guilds: DiscordGuild[]; expiresAt: number };
|
type CacheEntry = { guilds: DiscordGuild[]; expiresAt: number };
|
||||||
|
|
||||||
@@ -70,19 +70,19 @@ async function fetchUserGuildsUncached(accessToken: string): Promise<DiscordGuil
|
|||||||
const header = res.headers.get("retry-after");
|
const header = res.headers.get("retry-after");
|
||||||
retryAfter = header ? Number(header) : 1;
|
retryAfter = header ? Number(header) : 1;
|
||||||
}
|
}
|
||||||
const waitMs = Math.min(Math.max(retryAfter * 1000, 300), 5000);
|
const waitMs = Math.min(Math.max(retryAfter * 1000, 500), 15_000);
|
||||||
console.warn(`Discord guilds rate-limited — retry in ${waitMs}ms (attempt ${attempt}/${MAX_RETRIES})`);
|
console.warn(`Discord guilds rate-limited — retry in ${waitMs}ms (attempt ${attempt}/${MAX_RETRIES})`);
|
||||||
await new Promise((r) => setTimeout(r, waitMs));
|
await new Promise((r) => setTimeout(r, waitMs));
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (res.status === 401 || res.status === 403) {
|
||||||
|
throw new Error("Discord session expired — please sign in again.");
|
||||||
|
}
|
||||||
|
|
||||||
const body = await res.text().catch(() => "");
|
const body = await res.text().catch(() => "");
|
||||||
console.error("Discord guilds fetch failed:", res.status, body);
|
console.error("Discord guilds fetch failed:", res.status, body);
|
||||||
throw new Error(
|
throw new Error(`Failed to fetch your Discord servers. (${res.status})`);
|
||||||
res.status === 401
|
|
||||||
? "Discord session expired — please sign in again."
|
|
||||||
: `Failed to fetch your Discord servers. (${res.status})`
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
throw new Error("Failed to fetch your Discord servers. (rate limited)");
|
throw new Error("Failed to fetch your Discord servers. (rate limited)");
|
||||||
|
|||||||
Reference in New Issue
Block a user