Cleanup backend hotspots and apply code quality fixes

apps/server/convex/benchmarks.ts

@@ -1,6 +1,6 @@
+import type { FunctionReference } from "convex/server";
 import { v } from "convex/values";
 import { internalAction, internalMutation, query } from "./_generated/server";
-import { internal } from "./_generated/api";
 import { buildMatchingMap, type AAModel } from "./lib/model_matching";
 import { createLogger } from "./lib/logger";
 
@@ -33,6 +33,29 @@ const benchmarkValidator = v.object({
 	aime: v.optional(v.float64()),
 });
 
+const storeBenchmarksRef =
+	"benchmarks:storeBenchmarks" as unknown as FunctionReference<
+		"mutation",
+		"internal",
+		{
+			benchmarks: Array<{
+				openRouterModelId: string;
+				aaSlug: string;
+				aaCreatorName: string;
+				intelligenceIndex?: number;
+				codingIndex?: number;
+				mathIndex?: number;
+				mmluPro?: number;
+				gpqa?: number;
+				scicode?: number;
+				livecodebench?: number;
+				math500?: number;
+				aime?: number;
+			}>;
+		},
+		unknown
+	>;
+
 export const fetchAndStoreBenchmarks = internalAction({
 	args: {},
 	handler: async (ctx) => {
@@ -87,7 +110,7 @@ export const fetchAndStoreBenchmarks = internalAction({
 				}];
 			});
 
-			await ctx.runMutation(internal.benchmarks.storeBenchmarks, { benchmarks });
+			await ctx.runMutation(storeBenchmarksRef, { benchmarks });
 		} catch (error) {
 			void logger.error("Failed to refresh Artificial Analysis benchmarks", error);
 		}
@@ -184,4 +207,3 @@ export const getAllBenchmarks = query({
 	},
 });
 
-

apps/server/convex/chatTitle.ts

@@ -1,3 +1,4 @@
+import type { FunctionReference } from "convex/server";
 import type { Id } from "./_generated/dataModel";
 import type { ActionCtx } from "./_generated/server";
 import { action, internalAction, internalMutation, internalQuery, mutation } from "./_generated/server";
@@ -23,6 +24,38 @@ const TITLE_STYLE_PROMPTS: Record<"short" | "standard" | "long", string> = {
 type TitleLength = "short" | "standard" | "long";
 type TitleProvider = "osschat" | "openrouter";
 
+const getOpenRouterKeyInternalRef =
+	"users:getOpenRouterKeyInternal" as unknown as FunctionReference<
+		"query",
+		"internal",
+		{ userId: Id<"users"> },
+		string | null
+	>;
+
+const enforceTitleRateLimitRef =
+	"chatTitle:enforceTitleRateLimit" as unknown as FunctionReference<
+		"mutation",
+		"internal",
+		{ userId: Id<"users"> },
+		null
+	>;
+
+const getChatForTitleGenerationInternalRef =
+	"chatTitle:getChatForTitleGenerationInternal" as unknown as FunctionReference<
+		"query",
+		"internal",
+		{ chatId: Id<"chats">; userId: Id<"users"> },
+		{ title?: string | null } | null
+	>;
+
+const setGeneratedTitleInternalRef =
+	"chatTitle:setGeneratedTitleInternal" as unknown as FunctionReference<
+		"mutation",
+		"internal",
+		{ chatId: Id<"chats">; userId: Id<"users">; title: string; force?: boolean },
+		null
+	>;
+
 async function resolveOpenRouterKey(
 	ctx: ActionCtx,
 	userId: Id<"users">,
@@ -32,7 +65,7 @@ async function resolveOpenRouterKey(
 		return process.env.OPENROUTER_API_KEY ?? null;
 	}
 
-	const encryptedKey = await ctx.runQuery(internal.users.getOpenRouterKeyInternal, {
+	const encryptedKey = await ctx.runQuery(getOpenRouterKeyInternalRef, {
 		userId,
 	});
 	return encryptedKey ? await decryptSecret(encryptedKey) : null;
@@ -111,7 +144,7 @@ export const generateTitle = action({
 	returns: v.union(v.string(), v.null()),
 	handler: async (_ctx, args) => {
 		const userId = await requireAuthUserIdFromAction(_ctx, args.userId);
-		await _ctx.runMutation(internal.chatTitle.enforceTitleRateLimit, {
+		await _ctx.runMutation(enforceTitleRateLimitRef, {
 			userId,
 		});
 
@@ -140,7 +173,7 @@ export const generateAndSetTitleInternal = internalAction({
 		reason: v.optional(v.string()),
 	}),
 	handler: async (ctx, args) => {
-		const chat = await ctx.runQuery(internal.chatTitle.getChatForTitleGenerationInternal, {
+		const chat = await ctx.runQuery(getChatForTitleGenerationInternalRef, {
 			chatId: args.chatId,
 			userId: args.userId,
 		});
@@ -157,7 +190,7 @@ export const generateAndSetTitleInternal = internalAction({
 			return { saved: false, reason: "empty_seed" };
 		}
 
-		await ctx.runMutation(internal.chatTitle.enforceTitleRateLimit, {
+		await ctx.runMutation(enforceTitleRateLimitRef, {
 			userId: args.userId,
 		});
 
@@ -175,7 +208,7 @@ export const generateAndSetTitleInternal = internalAction({
 			return { saved: false, reason: "generation_failed" };
 		}
 
-		await ctx.runMutation(internal.chatTitle.setGeneratedTitleInternal, {
+		await ctx.runMutation(setGeneratedTitleInternalRef, {
 			chatId: args.chatId,
 			userId: args.userId,
 			title: generatedTitle,

apps/server/convex/chat_cleanup_helpers.ts

@@ -0,0 +1,80 @@
+import type { Id } from "./_generated/dataModel";
+import type { MutationCtx } from "./_generated/server";
+
+export async function completeActiveAndPendingStreams(
+	ctx: MutationCtx,
+	chatId: Id<"chats">,
+	now: number,
+) {
+	const [activeStreams, pendingStreams] = await Promise.all([
+		ctx.db
+			.query("streamJobs")
+			.withIndex("by_chat", (q) => q.eq("chatId", chatId).eq("status", "running"))
+			.collect(),
+		ctx.db
+			.query("streamJobs")
+			.withIndex("by_chat", (q) => q.eq("chatId", chatId).eq("status", "pending"))
+			.collect(),
+	]);
+
+	const streams = [...activeStreams, ...pendingStreams];
+	await Promise.all(
+		streams.map((stream) =>
+			ctx.db.patch(stream._id, {
+				status: "completed",
+				completedAt: now,
+			}),
+		),
+	);
+
+	return streams.length;
+}
+
+export async function softDeleteMessagesAfter(
+	ctx: MutationCtx,
+	chatId: Id<"chats">,
+	cutoffCreatedAt: number,
+	now: number,
+) {
+	const messagesToDelete = await ctx.db
+		.query("messages")
+		.withIndex("by_chat_not_deleted", (q) =>
+			q.eq("chatId", chatId).eq("deletedAt", undefined)
+		)
+		.order("asc")
+		.filter((q) => q.gt(q.field("createdAt"), cutoffCreatedAt))
+		.collect();
+
+	await Promise.all(
+		messagesToDelete.map((message) =>
+			ctx.db.patch(message._id, {
+				deletedAt: now,
+			}),
+		),
+	);
+
+	return messagesToDelete.length;
+}
+
+export async function softDeleteAllMessagesInChat(
+	ctx: MutationCtx,
+	chatId: Id<"chats">,
+	now: number,
+) {
+	const messages = await ctx.db
+		.query("messages")
+		.withIndex("by_chat_not_deleted", (q) =>
+			q.eq("chatId", chatId).eq("deletedAt", undefined)
+		)
+		.collect();
+
+	await Promise.all(
+		messages.map((message) =>
+			ctx.db.patch(message._id, {
+				deletedAt: now,
+			}),
+		),
+	);
+
+	return messages.length;
+}

apps/server/convex/chats.ts

@@ -7,6 +7,7 @@ import { rateLimiter } from "./lib/rateLimiter";
 import { throwRateLimitError } from "./lib/rateLimitUtils";
 import { sanitizeTitle } from "./lib/sanitize";
 import { requireAuthUserId } from "./lib/auth";
+import { softDeleteAllMessagesInChat } from "./chat_cleanup_helpers";
 
 const chatDoc = v.object({
 	_id: v.id("chats"),
@@ -165,34 +166,16 @@ export const remove = mutation({
 			throwRateLimitError("deletions", retryAfter);
 		}
 
-		const chat = await ctx.db.get(args.chatId);
-		if (!chat || chat.userId !== userId || chat.deletedAt) {
+		const chat = await assertOwnsChat(ctx, args.chatId, userId);
+		if (!chat) {
 			return { ok: false } as const;
 		}
 		const now = Date.now();
 
-		const messages = await ctx.db
-			.query("messages")
-			.withIndex("by_chat_not_deleted", (q) =>
-				q.eq("chatId", args.chatId).eq("deletedAt", undefined)
-			)
-			.collect();
-
-		await Promise.all(
-			messages.map((message) =>
-				ctx.db.patch(message._id, {
-					deletedAt: now,
-				}),
-			),
-		);
-
-		await ctx.db.patch(args.chatId, {
-			deletedAt: now,
-			messageCount: 0,
-		});
+		const deletedMessages = await softDeleteChatWithMessages(ctx, args.chatId, now);
 
 		await incrementStat(ctx, STAT_KEYS.CHATS_SOFT_DELETED);
-		await incrementStat(ctx, STAT_KEYS.MESSAGES_SOFT_DELETED, messages.length);
+		await incrementStat(ctx, STAT_KEYS.MESSAGES_SOFT_DELETED, deletedMessages);
 
 		return { ok: true } as const;
 	},
@@ -238,51 +221,23 @@ export const removeBulk = mutation({
 		let totalMessages = 0;
 
 		// First pass: validate all chats and collect valid ones
-		const validChats: Array<{ chatId: Id<"chats"> }> = [];
+		const validChats: Array<Id<"chats">> = [];
 		for (const chatId of args.chatIds) {
-			const chat = await ctx.db.get(chatId);
+			const chat = await assertOwnsChat(ctx, chatId, userId);
 
 			// Skip if chat doesn't exist, doesn't belong to user, or is already deleted
-			if (!chat || chat.userId !== userId || chat.deletedAt) {
+			if (!chat) {
 				failed++;
 				continue;
 			}
 
-			validChats.push({ chatId });
+			validChats.push(chatId);
 		}
 
-		// Second pass: fetch all messages for valid chats in parallel
-		const messagesByChat = await Promise.all(
-			validChats.map(async ({ chatId }) => {
-				const messages = await ctx.db
-					.query("messages")
-					.withIndex("by_chat_not_deleted", (q) =>
-						q.eq("chatId", chatId).eq("deletedAt", undefined)
-					)
-					.collect();
-				return { chatId, messages };
-			})
-		);
-
-		// Third pass: soft-delete all messages and chats
-		for (const { chatId, messages } of messagesByChat) {
-			// Soft-delete all messages for this chat
-			await Promise.all(
-				messages.map((message) =>
-					ctx.db.patch(message._id, {
-						deletedAt: now,
-					}),
-				),
-			);
-
-			// Soft-delete the chat
-			await ctx.db.patch(chatId, {
-				deletedAt: now,
-				messageCount: 0,
-			});
-
+		for (const chatId of validChats) {
+			const deletedMessages = await softDeleteChatWithMessages(ctx, chatId, now);
 			deleted++;
-			totalMessages += messages.length;
+			totalMessages += deletedMessages;
 		}
 
 		// Update stats
@@ -309,6 +264,19 @@ export async function assertOwnsChat(
 	return chat;
 }
 
+async function softDeleteChatWithMessages(
+	ctx: MutationCtx,
+	chatId: Id<"chats">,
+	now: number,
+) {
+	const deletedMessages = await softDeleteAllMessagesInChat(ctx, chatId, now);
+	await ctx.db.patch(chatId, {
+		deletedAt: now,
+		messageCount: 0,
+	});
+	return deletedMessages;
+}
+
 export const checkExportRateLimit = mutation({
 	args: {
 		userId: v.id("users"),

apps/server/convex/cleanupAction.ts

@@ -1,10 +1,27 @@
 "use node";
 
 import { createHmac, timingSafeEqual } from "node:crypto";
+import type { FunctionReference } from "convex/server";
 import { internalAction } from "./_generated/server";
-import { internal } from "./_generated/api";
 import { v } from "convex/values";
 
+const cleanupSoftDeletedRecordsRef =
+	"crons:cleanupSoftDeletedRecords" as unknown as FunctionReference<
+		"mutation",
+		"internal",
+		{
+			retentionDays?: number;
+			batchSize?: number;
+			dryRun?: boolean;
+		},
+		{
+			success: boolean;
+			deleted: number;
+			dryRun: boolean;
+			cutoffDate: string;
+		}
+	>;
+
 function safeCompare(a: string, b: string, hmacKey: string): boolean {
 	const hmacA = createHmac("sha256", hmacKey).update(a).digest();
 	const hmacB = createHmac("sha256", hmacKey).update(b).digest();
@@ -43,7 +60,7 @@ export const runCleanupBatchForWorkflow = internalAction({
 			throw new Error("batchSize must be between 1 and 1000");
 		}
 
-		const result = await ctx.runMutation(internal.crons.cleanupSoftDeletedRecords, {
+		const result = await ctx.runMutation(cleanupSoftDeletedRecordsRef, {
 			retentionDays,
 			batchSize,
 			dryRun: args.dryRun,