adjusts Agnai.chat and RisuAI rate limiting
This commit is contained in:
parent
4e05b01e90
commit
daf6a123d5
|
@ -21,7 +21,7 @@ kYuIJbnAGw5Oq0L6dXFW2DFwlcLz51kPVOmDc159FsQjyuPnta7NiZAANS8KM1CJ
|
|||
pwIDAQAB`;
|
||||
let IMPORTED_RISU_KEY: CryptoKey | null = null;
|
||||
|
||||
type RisuToken = { id: Uint8Array; expiresIn: number };
|
||||
type RisuToken = { id: string; expiresIn: number };
|
||||
type SignedToken = { data: RisuToken; sig: string };
|
||||
|
||||
(async () => {
|
||||
|
@ -54,14 +54,14 @@ export async function checkRisuToken(
|
|||
try {
|
||||
const { valid, data } = await validCheck(header);
|
||||
|
||||
if (!valid) {
|
||||
if (!valid || !data) {
|
||||
req.log.warn(
|
||||
{ token: header, data },
|
||||
"Invalid RisuAI token; using IP instead"
|
||||
);
|
||||
} else {
|
||||
req.log.info("RisuAI token validated");
|
||||
req.risuToken = header;
|
||||
req.risuToken = String(data.id);
|
||||
}
|
||||
} catch (err) {
|
||||
req.log.warn(
|
||||
|
@ -81,12 +81,13 @@ async function validCheck(header: string) {
|
|||
);
|
||||
} catch (err) {
|
||||
log.warn({ error: err.message }, "Provided unparseable RisuAI token");
|
||||
return { valid: false, data: "[unparseable]" };
|
||||
return { valid: false };
|
||||
}
|
||||
const data: RisuToken = tk.data;
|
||||
const sig = Buffer.from(tk.sig, "base64");
|
||||
|
||||
if (data.expiresIn < Math.floor(Date.now() / 1000)) {
|
||||
log.warn({ token: header }, "Provided expired RisuAI token");
|
||||
return { valid: false };
|
||||
}
|
||||
|
||||
|
@ -97,5 +98,9 @@ async function validCheck(header: string) {
|
|||
Buffer.from(JSON.stringify(data))
|
||||
);
|
||||
|
||||
if (!valid) {
|
||||
log.warn({ token: header }, "RisuAI token failed signature check");
|
||||
}
|
||||
|
||||
return { valid, data };
|
||||
}
|
||||
|
|
|
@ -32,7 +32,7 @@ const queue: Request[] = [];
|
|||
const log = logger.child({ module: "request-queue" });
|
||||
|
||||
/** Maximum number of queue slots for Agnai.chat requests. */
|
||||
const AGNAI_CONCURRENCY_LIMIT = 15;
|
||||
const AGNAI_CONCURRENCY_LIMIT = 5;
|
||||
/** Maximum number of queue slots for individual users. */
|
||||
const USER_CONCURRENCY_LIMIT = 1;
|
||||
|
||||
|
@ -68,12 +68,11 @@ export function enqueue(req: Request) {
|
|||
// more spots in the queue. Can't make it unlimited because people will
|
||||
// intentionally abuse it.
|
||||
// Authenticated users always get a single spot in the queue.
|
||||
const isAgnai = AGNAI_DOT_CHAT_IP.includes(req.ip);
|
||||
const maxConcurrentQueuedRequests =
|
||||
isGuest && req.ip === AGNAI_DOT_CHAT_IP
|
||||
? AGNAI_CONCURRENCY_LIMIT
|
||||
: USER_CONCURRENCY_LIMIT;
|
||||
isGuest && isAgnai ? AGNAI_CONCURRENCY_LIMIT : USER_CONCURRENCY_LIMIT;
|
||||
if (enqueuedRequestCount >= maxConcurrentQueuedRequests) {
|
||||
if (req.ip === AGNAI_DOT_CHAT_IP) {
|
||||
if (isAgnai) {
|
||||
// Re-enqueued requests are not counted towards the limit since they
|
||||
// already made it through the queue once.
|
||||
if (req.retryCount === 0) {
|
||||
|
@ -336,11 +335,7 @@ function killQueuedRequest(req: Request) {
|
|||
try {
|
||||
const message = `Your request has been terminated by the proxy because it has been in the queue for more than 5 minutes. The queue is currently ${queue.length} requests long.`;
|
||||
if (res.headersSent) {
|
||||
const fakeErrorEvent = buildFakeSse(
|
||||
"proxy queue error",
|
||||
message,
|
||||
req
|
||||
);
|
||||
const fakeErrorEvent = buildFakeSse("proxy queue error", message, req);
|
||||
res.write(fakeErrorEvent);
|
||||
res.end();
|
||||
} else {
|
||||
|
|
|
@ -1,7 +1,12 @@
|
|||
import { Request, Response, NextFunction } from "express";
|
||||
import { config } from "../config";
|
||||
|
||||
export const AGNAI_DOT_CHAT_IP = "157.230.249.32";
|
||||
export const AGNAI_DOT_CHAT_IP = [
|
||||
"157.230.249.32", // old
|
||||
"157.245.148.56",
|
||||
"174.138.29.50",
|
||||
"209.97.162.44",
|
||||
];
|
||||
|
||||
const RATE_LIMIT_ENABLED = Boolean(config.modelRateLimit);
|
||||
const RATE_LIMIT = Math.max(1, config.modelRateLimit);
|
||||
|
@ -66,7 +71,8 @@ export const ipLimiter = async (
|
|||
// Exempt Agnai.chat from rate limiting since it's shared between a lot of
|
||||
// users. Dunno how to prevent this from being abused without some sort of
|
||||
// identifier sent from Agnaistic to identify specific users.
|
||||
if (req.ip === AGNAI_DOT_CHAT_IP) {
|
||||
if (AGNAI_DOT_CHAT_IP.includes(req.ip)) {
|
||||
req.log.info("Exempting Agnai request from rate limiting.");
|
||||
next();
|
||||
return;
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue