feat(migrations): clean up empty conversations from db (#1561)

* feat(migrations): clean up empty conversations from db

* fix(tests): better test coverage, fix bug with sessions

* feat: better migration

* fix: stop test server after running tests

* fix: migration code

* fix: refactor to take into account large size of session collection

scripts/setupTest.ts

@@ -1,4 +1,4 @@
-import { vi } from "vitest";
+import { vi, afterAll } from "vitest";
 import dotenv from "dotenv";
 import { resolve } from "path";
 import fs from "fs";
@@ -41,3 +41,9 @@ vi.mock("$env/dynamic/private", async () => {
 		},
 	};
 });
+
+afterAll(async () => {
+	if (mongoServer) {
+		await mongoServer.stop();
+	}
+});

src/lib/migrations/routines/09-delete-empty-conversations.spec.ts

@@ -0,0 +1,214 @@
+import type { Session } from "$lib/types/Session";
+import type { User } from "$lib/types/User";
+import type { Conversation } from "$lib/types/Conversation";
+import { ObjectId } from "mongodb";
+import { deleteConversations } from "./09-delete-empty-conversations";
+import { afterAll, afterEach, beforeAll, describe, expect, test } from "vitest";
+import { collections } from "$lib/server/database";
+
+type Message = Conversation["messages"][number];
+
+const userData = {
+	_id: new ObjectId(),
+	createdAt: new Date(),
+	updatedAt: new Date(),
+	username: "new-username",
+	name: "name",
+	avatarUrl: "https://example.com/avatar.png",
+	hfUserId: "9999999999",
+} satisfies User;
+Object.freeze(userData);
+
+const sessionForUser = {
+	_id: new ObjectId(),
+	createdAt: new Date(),
+	updatedAt: new Date(),
+	userId: userData._id,
+	sessionId: "session-id-9999999999",
+	expiresAt: new Date(Date.now() + 1000 * 60 * 60 * 24),
+} satisfies Session;
+Object.freeze(sessionForUser);
+
+const userMessage = {
+	from: "user",
+	id: "user-message-id",
+	content: "Hello, how are you?",
+} satisfies Message;
+
+const assistantMessage = {
+	from: "assistant",
+	id: "assistant-message-id",
+	content: "I'm fine, thank you!",
+} satisfies Message;
+
+const systemMessage = {
+	from: "system",
+	id: "system-message-id",
+	content: "This is a system message",
+} satisfies Message;
+
+const conversationBase = {
+	_id: new ObjectId(),
+	createdAt: new Date(Date.now() - 7 * 24 * 60 * 60 * 1000),
+	updatedAt: new Date(Date.now() - 7 * 24 * 60 * 60 * 1000),
+	model: "model-id",
+	embeddingModel: "embedding-model-id",
+	title: "title",
+	messages: [],
+} satisfies Conversation;
+
+describe.sequential("Deleting discarded conversations", async () => {
+	test("a conversation with no messages should get deleted", async () => {
+		await collections.conversations.insertOne({
+			...conversationBase,
+			sessionId: sessionForUser.sessionId,
+		});
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(1);
+	});
+	test("a conversation with no messages that is less than 1 hour old should not get deleted", async () => {
+		await collections.conversations.insertOne({
+			...conversationBase,
+			sessionId: sessionForUser.sessionId,
+			createdAt: new Date(Date.now() - 30 * 60 * 1000),
+		});
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(0);
+	});
+	test("a conversation with only system messages should get deleted", async () => {
+		await collections.conversations.insertOne({
+			...conversationBase,
+			sessionId: sessionForUser.sessionId,
+			messages: [systemMessage],
+		});
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(1);
+	});
+	test("a conversation with a user message should not get deleted", async () => {
+		await collections.conversations.insertOne({
+			...conversationBase,
+			sessionId: sessionForUser.sessionId,
+			messages: [userMessage],
+		});
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(0);
+	});
+	test("a conversation with an assistant message should not get deleted", async () => {
+		await collections.conversations.insertOne({
+			...conversationBase,
+			sessionId: sessionForUser.sessionId,
+			messages: [assistantMessage],
+		});
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(0);
+	});
+	test("a conversation with a mix of messages should not get deleted", async () => {
+		await collections.conversations.insertOne({
+			...conversationBase,
+			sessionId: sessionForUser.sessionId,
+			messages: [systemMessage, userMessage, assistantMessage, userMessage, assistantMessage],
+		});
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(0);
+	});
+	test("a conversation with a userId and no sessionId should not get deleted", async () => {
+		await collections.conversations.insertOne({
+			...conversationBase,
+			messages: [userMessage, assistantMessage],
+			userId: userData._id,
+		});
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(0);
+	});
+	test("a conversation with no userId or sessionId should get deleted", async () => {
+		await collections.conversations.insertOne({
+			...conversationBase,
+			messages: [userMessage, assistantMessage],
+		});
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(1);
+	});
+	test("a conversation with a sessionId that exists should not get deleted", async () => {
+		await collections.conversations.insertOne({
+			...conversationBase,
+			messages: [userMessage, assistantMessage],
+			sessionId: sessionForUser.sessionId,
+		});
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(0);
+	});
+	test("a conversation with a userId and a sessionId that doesn't exist should NOT get deleted", async () => {
+		await collections.conversations.insertOne({
+			...conversationBase,
+			userId: userData._id,
+			messages: [userMessage, assistantMessage],
+			sessionId: new ObjectId().toString(),
+		});
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(0);
+	});
+	test("a conversation with only a sessionId that doesn't exist, should get deleted", async () => {
+		await collections.conversations.insertOne({
+			...conversationBase,
+			messages: [userMessage, assistantMessage],
+			sessionId: new ObjectId().toString(),
+		});
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(1);
+	});
+	test("many conversations should get deleted", async () => {
+		const conversations = Array.from({ length: 10010 }, () => ({
+			...conversationBase,
+			_id: new ObjectId(),
+		}));
+
+		await collections.conversations.insertMany(conversations);
+
+		const result = await deleteConversations(collections);
+
+		expect(result).toBe(10010);
+	});
+});
+
+beforeAll(async () => {
+	await collections.users.insertOne(userData);
+	await collections.sessions.insertOne(sessionForUser);
+});
+
+afterAll(async () => {
+	await collections.users.deleteOne({
+		_id: userData._id,
+	});
+	await collections.sessions.deleteOne({
+		_id: sessionForUser._id,
+	});
+	await collections.conversations.deleteMany({});
+});
+
+afterEach(async () => {
+	await collections.conversations.deleteMany({
+		_id: { $in: [conversationBase._id] },
+	});
+});

src/lib/migrations/routines/09-delete-empty-conversations.ts

@@ -0,0 +1,88 @@
+import type { Migration } from ".";
+import { collections } from "$lib/server/database";
+import { Collection, FindCursor, ObjectId } from "mongodb";
+import { logger } from "$lib/server/logger";
+import type { Conversation } from "$lib/types/Conversation";
+
+const BATCH_SIZE = 1000;
+const DELETE_THRESHOLD_MS = 60 * 60 * 1000;
+
+async function deleteBatch(conversations: Collection<Conversation>, ids: ObjectId[]) {
+	if (ids.length === 0) return 0;
+	const deleteResult = await conversations.deleteMany({ _id: { $in: ids } });
+	return deleteResult.deletedCount;
+}
+
+async function processCursor<T>(
+	cursor: FindCursor<T>,
+	processBatchFn: (batch: T[]) => Promise<void>
+) {
+	let batch = [];
+	while (await cursor.hasNext()) {
+		const doc = await cursor.next();
+		if (doc) {
+			batch.push(doc);
+		}
+		if (batch.length >= BATCH_SIZE) {
+			await processBatchFn(batch);
+			batch = [];
+		}
+	}
+	if (batch.length > 0) {
+		await processBatchFn(batch);
+	}
+}
+
+export async function deleteConversations(
+	collections: typeof import("$lib/server/database").collections
+) {
+	let deleteCount = 0;
+	const { conversations, sessions } = collections;
+
+	// First criteria: Delete conversations with no user/assistant messages older than 1 hour
+	const emptyConvCursor = conversations
+		.find({
+			"messages.from": { $not: { $in: ["user", "assistant"] } },
+			createdAt: { $lt: new Date(Date.now() - DELETE_THRESHOLD_MS) },
+		})
+		.batchSize(BATCH_SIZE);
+
+	await processCursor(emptyConvCursor, async (batch) => {
+		const ids = batch.map((doc) => doc._id);
+		deleteCount += await deleteBatch(conversations, ids);
+	});
+
+	// Second criteria: Process conversations without users in batches and check sessions
+	const noUserCursor = conversations.find({ userId: { $exists: false } }).batchSize(BATCH_SIZE);
+
+	await processCursor(noUserCursor, async (batch) => {
+		const sessionIds = [
+			...new Set(batch.map((conv) => conv.sessionId).filter((id): id is string => !!id)),
+		];
+
+		const existingSessions = await sessions.find({ sessionId: { $in: sessionIds } }).toArray();
+		const validSessionIds = new Set(existingSessions.map((s) => s.sessionId));
+
+		const invalidConvs = batch.filter(
+			(conv) => !conv.sessionId || !validSessionIds.has(conv.sessionId)
+		);
+		const idsToDelete = invalidConvs.map((conv) => conv._id);
+		deleteCount += await deleteBatch(conversations, idsToDelete);
+	});
+
+	logger.info(`[MIGRATIONS] Deleted ${deleteCount} conversations in total.`);
+	return deleteCount;
+}
+
+const deleteEmptyConversations: Migration = {
+	_id: new ObjectId("000000000000000000000009"),
+	name: "Delete conversations with no user or assistant messages or valid sessions",
+	up: async () => {
+		await deleteConversations(collections);
+		return true;
+	},
+	runEveryTime: true,
+	runForHuggingChat: "only",
+};
+
+export default deleteEmptyConversations;

src/lib/migrations/routines/index.ts

@@ -9,7 +9,7 @@ import updateMessageFiles from "./05-update-message-files";
 import trimMessageUpdates from "./06-trim-message-updates";
 import resetTools from "./07-reset-tools-in-settings";
 import updateFeaturedToReview from "./08-update-featured-to-review";
-
+import deleteEmptyConversations from "./09-delete-empty-conversations";
 export interface Migration {
 	_id: ObjectId;
 	name: string;
@@ -29,4 +29,5 @@ export const migrations: Migration[] = [
 	trimMessageUpdates,
 	resetTools,
 	updateFeaturedToReview,
+	deleteEmptyConversations,
 ];

src/lib/server/database.ts

@@ -228,6 +228,20 @@ export class Database {
 		tools.createIndex({ createdById: 1, userCount: -1 }).catch((e) => logger.error(e));
 		tools.createIndex({ userCount: 1 }).catch((e) => logger.error(e));
 		tools.createIndex({ last24HoursCount: 1 }).catch((e) => logger.error(e));
+
+		conversations
+			.createIndex({
+				"messages.from": 1,
+				createdAt: 1,
+			})
+			.catch((e) => logger.error(e));
+
+		conversations
+			.createIndex({
+				userId: 1,
+				sessionId: 1,
+			})
+			.catch((e) => logger.error(e));
 	}
 }
 

vite.config.ts

@@ -41,5 +41,6 @@ export default defineConfig({
 		setupFiles: ["./scripts/setupTest.ts"],
 		deps: { inline: ["@sveltejs/kit"] },
 		globals: true,
+		testTimeout: 10000,
 	},
 });