very basic rate limiter (#320)

.env

@@ -70,6 +70,7 @@ PARQUET_EXPORT_DATASET=
 PARQUET_EXPORT_HF_TOKEN=
 PARQUET_EXPORT_SECRET=
 
+RATE_LIMIT= # requests per minute
 
 PUBLIC_APP_NAME=ChatUI # name used as title throughout the app
 PUBLIC_APP_ASSETS=chatui # used to find logos & favicons in static/$PUBLIC_APP_ASSETS

src/lib/server/database.ts

@@ -6,6 +6,7 @@ import type { WebSearch } from "$lib/types/WebSearch";
 import type { AbortedGeneration } from "$lib/types/AbortedGeneration";
 import type { Settings } from "$lib/types/Settings";
 import type { User } from "$lib/types/User";
+import type { MessageEvent } from "$lib/types/MessageEvent";
 
 if (!MONGODB_URL) {
 	throw new Error(
@@ -27,6 +28,7 @@ const abortedGenerations = db.collection<AbortedGeneration>("abortedGenerations"
 const settings = db.collection<Settings>("settings");
 const users = db.collection<User>("users");
 const webSearches = db.collection<WebSearch>("webSearches");
+const messageEvents = db.collection<MessageEvent>("messageEvents");
 
 export { client, db };
 export const collections = {
@@ -36,6 +38,7 @@ export const collections = {
 	settings,
 	users,
 	webSearches,
+	messageEvents,
 };
 
 client.on("open", () => {
@@ -59,4 +62,5 @@ client.on("open", () => {
 	settings.createIndex({ userId: 1 }, { unique: true, sparse: true }).catch(console.error);
 	users.createIndex({ hfUserId: 1 }, { unique: true }).catch(console.error);
 	users.createIndex({ sessionId: 1 }, { unique: true, sparse: true }).catch(console.error);
+	messageEvents.createIndex({ createdAt: 1 }, { expireAfterSeconds: 60 }).catch(console.error);
 });

src/lib/stores/errors.ts

@@ -3,6 +3,7 @@ import { writable } from "svelte/store";
 export const ERROR_MESSAGES = {
 	default: "Oops, something went wrong.",
 	authOnly: "You have to be logged in.",
+	rateLimited: "You are sending too many messages. Try again later.",
 };
 
 export const error = writable<string | null>(null);

src/lib/types/MessageEvent.ts

@@ -0,0 +1,6 @@
+import type { Timestamps } from "./Timestamps";
+import type { User } from "./User";
+
+export interface MessageEvent extends Pick<Timestamps, "createdAt"> {
+	userId: User["_id"] | User["sessionId"];
+}

src/routes/conversation/[id]/+page.svelte

@@ -207,6 +207,8 @@
 		} catch (err) {
 			if (err instanceof Error && err.message.includes("overloaded")) {
 				$error = "Too much traffic, please try again.";
+			} else if (err instanceof Error && err.message.includes("429")) {
+				$error = ERROR_MESSAGES.rateLimited;
 			} else if (err instanceof Error) {
 				$error = err.message;
 			} else {

src/routes/conversation/[id]/+server.ts

@@ -1,3 +1,4 @@
+import { RATE_LIMIT } from "$env/static/private";
 import { buildPrompt } from "$lib/buildPrompt";
 import { PUBLIC_SEP_TOKEN } from "$lib/constants/publicSepToken";
 import { abortedGenerations } from "$lib/server/abortedGenerations";
@@ -5,6 +6,7 @@ import { authCondition } from "$lib/server/auth";
 import { collections } from "$lib/server/database";
 import { modelEndpoint } from "$lib/server/modelEndpoint";
 import { models } from "$lib/server/models";
+import { ERROR_MESSAGES } from "$lib/stores/errors.js";
 import type { Message } from "$lib/types/Message";
 import { concatUint8Arrays } from "$lib/utils/concatUint8Arrays";
 import { streamToAsyncIterable } from "$lib/utils/streamToAsyncIterable";
@@ -20,6 +22,12 @@ export async function POST({ request, fetch, locals, params }) {
 	const convId = new ObjectId(id);
 	const date = new Date();
 
+	const userId = locals.user?._id ?? locals.sessionId;
+
+	if (!userId) {
+		throw error(401, "Unauthorized");
+	}
+
 	const conv = await collections.conversations.findOne({
 		_id: convId,
 		...authCondition(locals),
@@ -29,6 +37,12 @@ export async function POST({ request, fetch, locals, params }) {
 		throw error(404, "Conversation not found");
 	}
 
+	const nEvents = await collections.messageEvents.countDocuments({ userId });
+
+	if (RATE_LIMIT != "" && nEvents > parseInt(RATE_LIMIT)) {
+		throw error(429, ERROR_MESSAGES.rateLimited);
+	}
+
 	const model = models.find((m) => m.id === conv.model);
 
 	if (!model) {
@@ -118,6 +132,11 @@ export async function POST({ request, fetch, locals, params }) {
 			id: (responseId as Message["id"]) || crypto.randomUUID(),
 		});
 
+		await collections.messageEvents.insertOne({
+			userId: userId,
+			createdAt: new Date(),
+		});
+
 		await collections.conversations.updateOne(
 			{
 				_id: convId,