src/lib/server/database.ts

import { GridFSBucket, MongoClient, ReadPreference } from "mongodb";
import type { Conversation } from "$lib/types/Conversation";
import type { SharedConversation } from "$lib/types/SharedConversation";
import type { AbortedGeneration } from "$lib/types/AbortedGeneration";
import type { Generation, GenerationEvent } from "$lib/types/Generation";
import type { Settings } from "$lib/types/Settings";
import type { User } from "$lib/types/User";
import type { MessageEvent } from "$lib/types/MessageEvent";
import type { Session } from "$lib/types/Session";
import type { Assistant } from "$lib/types/Assistant";
import type { Report } from "$lib/types/Report";
import type { ConversationStats } from "$lib/types/ConversationStats";
import type { MigrationResult } from "$lib/types/MigrationResult";
import type { Semaphore } from "$lib/types/Semaphore";
import { MongoMemoryServer } from "mongodb-memory-server";
import { logger } from "$lib/server/logger";
import { building } from "$app/environment";
import type { TokenCache } from "$lib/types/TokenCache";
import { onExit } from "./exitHandler";
import { fileURLToPath } from "url";
import { dirname, join } from "path";
import { existsSync, mkdirSync } from "fs";
import { findRepoRoot } from "./findRepoRoot";
import type { ConfigKey } from "$lib/types/ConfigKey";
import { config } from "$lib/server/config";
import type {
	AgentConfigurationLedger,
	AgentConfigurationReceipt,
	Workspace,
} from "$lib/types/Workspace";

export const CONVERSATION_STATS_COLLECTION = "conversations.stats";

export class Database {
	private client?: MongoClient;
	private mongoServer?: MongoMemoryServer;

	private static instance: Database;

	private async init() {
		const DB_FOLDER =
			config.MONGO_STORAGE_PATH ||
			join(findRepoRoot(dirname(fileURLToPath(import.meta.url))), "db");

		if (!config.MONGODB_URL) {
			logger.warn("No MongoDB URL found, using in-memory server");

			logger.info(`Using database path: ${DB_FOLDER}`);
			// Create db directory if it doesn't exist
			if (!existsSync(DB_FOLDER)) {
				logger.info(`Creating database directory at ${DB_FOLDER}`);
				mkdirSync(DB_FOLDER, { recursive: true });
			}

			this.mongoServer = await MongoMemoryServer.create({
				instance: {
					dbName: config.MONGODB_DB_NAME + (import.meta.env.MODE === "test" ? "-test" : ""),
					dbPath: DB_FOLDER,
				},
				binary: {
					version: "7.0.18",
				},
			});
			this.client = new MongoClient(this.mongoServer.getUri(), {
				directConnection: config.MONGODB_DIRECT_CONNECTION === "true",
			});
		} else {
			this.client = new MongoClient(config.MONGODB_URL, {
				directConnection: config.MONGODB_DIRECT_CONNECTION === "true",
			});
		}

		try {
			logger.info("Connecting to database");
			await this.client.connect();
			logger.info("Connected to database");
			this.client.db(config.MONGODB_DB_NAME + (import.meta.env.MODE === "test" ? "-test" : ""));
			await this.initDatabase();
		} catch (err) {
			logger.error(err, "Error connecting to database");
			process.exit(1);
		}

		// Disconnect DB on exit. Registered `last` so other exit handlers (e.g. the
		// reaper finalizing in-flight generations) finish their writes before the
		// client is force-closed.
		onExit(
			async () => {
				logger.info("Closing database connection");
				await this.client?.close(true);
				await this.mongoServer?.stop();
			},
			{ last: true }
		);
	}

	public static async getInstance(): Promise<Database> {
		if (!Database.instance) {
			Database.instance = new Database();
			await Database.instance.init();
		}

		return Database.instance;
	}

	/**
	 * Return mongoClient
	 */
	public getClient(): MongoClient {
		if (!this.client) {
			throw new Error("Database not initialized");
		}

		return this.client;
	}

	/**
	 * Return map of database's collections
	 */
	public getCollections() {
		if (!this.client) {
			throw new Error("Database not initialized");
		}

		const db = this.client.db(
			config.MONGODB_DB_NAME + (import.meta.env.MODE === "test" ? "-test" : "")
		);

		// Collections with default readPreference (primary) - critical for read-after-write consistency
		const conversations = db.collection<Conversation>("conversations");
		const workspaces = db.collection<Workspace>("workspaces");
		const agentConfigurationLedgers = db.collection<AgentConfigurationLedger>(
			"agentConfigurationLedgers"
		);
		const agentConfigurationReceipts = db.collection<AgentConfigurationReceipt>(
			"agentConfigurationReceipts"
		);
		const settings = db.collection<Settings>("settings");
		const users = db.collection<User>("users");
		const sessions = db.collection<Session>("sessions");
		const messageEvents = db.collection<MessageEvent>("messageEvents");
		const abortedGenerations = db.collection<AbortedGeneration>("abortedGenerations");
		const generations = db.collection<Generation>("generations");
		const generationEvents = db.collection<GenerationEvent>("generationEvents");
		const semaphores = db.collection<Semaphore>("semaphores");
		const tokenCaches = db.collection<TokenCache>("tokens");
		const configCollection = db.collection<ConfigKey>("config");
		const migrationResults = db.collection<MigrationResult>("migrationResults");
		const sharedConversations = db.collection<SharedConversation>("sharedConversations");
		const bucket = new GridFSBucket(db, { bucketName: "files" });

		// Collections with secondaryPreferred - heavy reads, can tolerate slight replication lag
		const secondaryPreferred = ReadPreference.SECONDARY_PREFERRED;
		const assistants = db.collection<Assistant>("assistants", {
			readPreference: secondaryPreferred,
		});
		const conversationStats = db.collection<ConversationStats>(CONVERSATION_STATS_COLLECTION, {
			readPreference: secondaryPreferred,
		});
		const reports = db.collection<Report>("reports", {
			readPreference: secondaryPreferred,
		});
		const tools = db.collection("tools", {
			readPreference: secondaryPreferred,
		});
		return {
			conversations,
			workspaces,
			agentConfigurationLedgers,
			agentConfigurationReceipts,
			conversationStats,
			assistants,
			reports,
			sharedConversations,
			abortedGenerations,
			generations,
			generationEvents,
			settings,
			users,
			sessions,
			messageEvents,
			bucket,
			migrationResults,
			semaphores,
			tokenCaches,
			tools,
			config: configCollection,
		};
	}

	/**
	 * Init database once connected: Index creation
	 * @private
	 */
	private async initDatabase() {
		const {
			conversations,
			workspaces,
			agentConfigurationLedgers,
			agentConfigurationReceipts,
			conversationStats,
			assistants,
			reports,
			sharedConversations,
			abortedGenerations,
			generations,
			generationEvents,
			settings,
			users,
			sessions,
			messageEvents,
			semaphores,
			tokenCaches,
			config,
		} = this.getCollections();

		// A conversation owns at most one durable agent workspace. Owner/status
		// indexes keep authorization and lifecycle sweeps bounded as the table grows.
		workspaces
			.createIndex({ conversationId: 1 }, { unique: true })
			.catch((e) => logger.error(e, "Error creating unique workspace conversation index"));
		workspaces
			.createIndex(
				{ userId: 1, "sandbox.status": 1, updatedAt: -1 },
				{ partialFilterExpression: { userId: { $exists: true } } }
			)
			.catch((e) => logger.error(e, "Error creating workspace owner/status index"));
		workspaces
			.createIndex(
				{ sessionId: 1, "sandbox.status": 1, updatedAt: -1 },
				{ partialFilterExpression: { sessionId: { $exists: true } } }
			)
			.catch((e) => logger.error(e, "Error creating workspace session/status index"));
		workspaces
			.createIndex(
				{ "sandbox.id": 1 },
				{
					unique: true,
					partialFilterExpression: { "sandbox.id": { $type: "string" } },
				}
			)
			.catch((e) => logger.error(e, "Error creating unique workspace sandbox job index"));
		workspaces
			.createIndex({ "sandbox.status": 1, "sandbox.lastSeenAt": 1 })
			.catch((e) => logger.error(e, "Error creating workspace lifecycle index"));
		workspaces
			.createIndex(
				{ "llmProxy.tokenHash": 1 },
				{
					unique: true,
					partialFilterExpression: { "llmProxy.tokenHash": { $type: "string" } },
				}
			)
			.catch((e) => logger.error(e, "Error creating unique workspace LLM proxy index"));
		// A browser capability must resolve to at most one exact workspace. Startup
		// fails if this security boundary cannot be installed.
		await workspaces.createIndex(
			{ "preview.capabilityHash": 1 },
			{
				name: "unique_workspace_preview_capability",
				unique: true,
				partialFilterExpression: { "preview.capabilityHash": { $type: "string" } },
			}
		);
		// The document id is the Workspace id. This secondary key makes orphan
		// cleanup by conversation bounded without weakening the exactly-once log.
		agentConfigurationLedgers
			.createIndex({ conversationId: 1 })
			.catch((e) => logger.error(e, "Error creating Agent configuration ledger index"));
		await agentConfigurationReceipts.createIndex(
			{ workspaceId: 1, controlId: 1 },
			{ name: "unique_agent_configuration_control", unique: true }
		);
		await agentConfigurationReceipts.createIndex(
			{ workspaceId: 1, sequence: 1 },
			{
				name: "unique_agent_configuration_sequence",
				unique: true,
				partialFilterExpression: { sequence: { $type: "number" } },
			}
		);
		agentConfigurationReceipts
			.createIndex({ conversationId: 1 })
			.catch((e) => logger.error(e, "Error creating Agent configuration receipt index"));

		conversations
			.createIndex(
				{ sessionId: 1, updatedAt: -1 },
				{ partialFilterExpression: { sessionId: { $exists: true } } }
			)
			.catch((e) =>
				logger.error(e, "Error creating index for conversations by sessionId and updatedAt")
			);
		conversations
			.createIndex(
				{ userId: 1, updatedAt: -1 },
				{ partialFilterExpression: { userId: { $exists: true } } }
			)
			.catch((e) =>
				logger.error(e, "Error creating index for conversations by userId and updatedAt")
			);
		conversations
			.createIndex(
				{ "message.id": 1, "message.ancestors": 1 },
				{ partialFilterExpression: { userId: { $exists: true } } }
			)
			.catch((e) =>
				logger.error(e, "Error creating index for conversations by messageId and ancestors")
			);
		// Not strictly necessary, could use _id, but more convenient. Also for stats
		// To do stats on conversation messages
		conversations
			.createIndex({ "messages.createdAt": 1 }, { sparse: true })
			.catch((e) =>
				logger.error(e, "Error creating index for conversations by messages createdAt")
			);
		// Unique index for stats
		conversationStats
			.createIndex(
				{
					type: 1,
					"date.field": 1,
					"date.span": 1,
					"date.at": 1,
					distinct: 1,
				},
				{ unique: true }
			)
			.catch((e) =>
				logger.error(
					e,
					"Error creating index for conversationStats by type, date.field and date.span"
				)
			);
		// Allow easy check of last computed stat for given type/dateField
		conversationStats
			.createIndex({
				type: 1,
				"date.field": 1,
				"date.at": 1,
			})
			.catch((e) => logger.error(e, "Error creating index for abortedGenerations by updatedAt"));
		abortedGenerations
			.createIndex({ updatedAt: 1 }, { expireAfterSeconds: 30 })
			.catch((e) =>
				logger.error(
					e,
					"Error creating index for abortedGenerations by updatedAt and expireAfterSeconds"
				)
			);
		// Stop markers are generation-scoped. Replace the former conversation-only
		// unique index so a late marker from generation A cannot block or overwrite
		// generation B's marker in the same conversation.
		try {
			await abortedGenerations.dropIndex("conversationId_1");
		} catch (cause) {
			if (![26, 27].includes((cause as { code?: number }).code ?? -1)) throw cause;
		}
		await abortedGenerations.createIndex(
			{ conversationId: 1, generationId: 1 },
			{
				name: "conversation_generation_unique",
				unique: true,
				partialFilterExpression: { generationId: { $type: "string" } },
			}
		);
		// This is the durable workspace execution claim. The bridge event stream is
		// workspace-global, so two running Agent generations could otherwise mirror
		// and terminate each other's turns. Unlike best-effort indexes below, startup
		// must fail if this invariant cannot be installed.
		await generations.createIndex(
			{ workspaceId: 1 },
			{
				name: "one_running_agent_generation_per_workspace",
				unique: true,
				partialFilterExpression: {
					mode: "agent",
					status: "running",
					workspaceId: { $type: "objectId" },
				},
			}
		);
		// Client generation ids are idempotency keys, not advisory metadata. The
		// route reserves one before appending a message, so startup must fail closed
		// if MongoDB cannot enforce the corresponding global uniqueness invariant.
		await generations.createIndex({ generationId: 1 }, { unique: true });
		generations
			.createIndex({ conversationId: 1, startedAt: -1 })
			.catch((e) =>
				logger.error(e, "Error creating index for generations by conversationId and startedAt")
			);
		generations
			.createIndex(
				{ userId: 1, updatedAt: -1 },
				{ partialFilterExpression: { userId: { $exists: true } } }
			)
			.catch((e) => logger.error(e, "Error creating index for generations by userId"));
		generations
			.createIndex(
				{ sessionId: 1, updatedAt: -1 },
				{ partialFilterExpression: { sessionId: { $exists: true } } }
			)
			.catch((e) => logger.error(e, "Error creating index for generations by sessionId"));
		// Serves the reaper's query: still-running runs ordered by heartbeat age.
		generations
			.createIndex({ status: 1, lastHeartbeatAt: 1 })
			.catch((e) =>
				logger.error(e, "Error creating index for generations by status and lastHeartbeatAt")
			);
		// Terminal recovery markers deliberately survive the status CAS until
		// message/workspace side effects finish, so the reaper must also find them
		// efficiently on already-terminal rows.
		generations
			.createIndex(
				{ "agent.terminalization.expiresAt": 1 },
				{
					name: "agent_terminalization_expiry",
					partialFilterExpression: {
						"agent.terminalization.expiresAt": { $type: "date" },
					},
				}
			)
			.catch((e) => logger.error(e, "Error creating index for Agent terminal recovery"));
		generations
			.createIndex(
				{ endedAt: 1 },
				{
					expireAfterSeconds: 7 * 24 * 60 * 60,
					partialFilterExpression: { endedAt: { $exists: true } },
				}
			)
			.catch((e) => logger.error(e, "Error creating TTL index for generations by endedAt"));

		// Unique so a retried insert is idempotent, and compound so it is also the
		// exact index the replay/tail range scan uses.
		generationEvents
			.createIndex({ generationId: 1, seq: 1 }, { unique: true })
			.catch((e) =>
				logger.error(e, "Error creating index for generationEvents by generationId and seq")
			);
		generationEvents
			.createIndex(
				{ generationId: 1, sourceSequence: 1, sourceIndex: 1 },
				{
					unique: true,
					partialFilterExpression: { sourceSequence: { $exists: true } },
				}
			)
			.catch((e) => logger.error(e, "Error creating Agent bridge event idempotency index"));
		generationEvents
			.createIndex(
				{ generationId: 1, systemKey: 1, systemIndex: 1 },
				{
					unique: true,
					partialFilterExpression: { systemKey: { $exists: true } },
				}
			)
			.catch((e) => logger.error(e, "Error creating Agent system event idempotency index"));
		generationEvents
			.createIndex({ createdAt: 1 }, { expireAfterSeconds: 24 * 60 * 60 })
			.catch((e) => logger.error(e, "Error creating TTL index for generationEvents by createdAt"));

		sharedConversations.createIndex({ hash: 1 }, { unique: true }).catch((e) => logger.error(e));
		settings
			.createIndex({ sessionId: 1 }, { unique: true, sparse: true })
			.catch((e) => logger.error(e, "Error creating index for settings by sessionId"));
		settings
			.createIndex({ userId: 1 }, { unique: true, sparse: true })
			.catch((e) => logger.error(e, "Error creating index for settings by userId"));
		settings
			.createIndex({ assistants: 1 })
			.catch((e) => logger.error(e, "Error creating index for settings by assistants"));
		users
			.createIndex({ hfUserId: 1 }, { unique: true })
			.catch((e) => logger.error(e, "Error creating index for users by hfUserId"));
		users
			.createIndex({ sessionId: 1 }, { unique: true, sparse: true })
			.catch((e) => logger.error(e, "Error creating index for users by sessionId"));
		// No unicity because due to renames & outdated info from oauth provider, there may be the same username on different users
		users
			.createIndex({ username: 1 })
			.catch((e) => logger.error(e, "Error creating index for users by username"));
		// For stats queries filtering users by creation date
		users
			.createIndex({ createdAt: 1 })
			.catch((e) => logger.error(e, "Error creating index for users by createdAt"));
		messageEvents
			.createIndex({ expiresAt: 1 }, { expireAfterSeconds: 1 })
			.catch((e) => logger.error(e, "Error creating index for messageEvents by expiresAt"));
		sessions.createIndex({ expiresAt: 1 }, { expireAfterSeconds: 0 }).catch((e) => logger.error(e));
		sessions
			.createIndex({ sessionId: 1 }, { unique: true })
			.catch((e) => logger.error(e, "Error creating index for sessions by sessionId"));
		assistants
			.createIndex({ createdById: 1, userCount: -1 })
			.catch((e) =>
				logger.error(e, "Error creating index for assistants by createdById and userCount")
			);
		assistants
			.createIndex({ userCount: 1 })
			.catch((e) => logger.error(e, "Error creating index for assistants by userCount"));
		assistants
			.createIndex({ review: 1, userCount: -1 })
			.catch((e) => logger.error(e, "Error creating index for assistants by review and userCount"));
		assistants
			.createIndex({ modelId: 1, userCount: -1 })
			.catch((e) =>
				logger.error(e, "Error creating index for assistants by modelId and userCount")
			);
		assistants
			.createIndex({ searchTokens: 1 })
			.catch((e) => logger.error(e, "Error creating index for assistants by searchTokens"));
		assistants
			.createIndex({ last24HoursCount: 1 })
			.catch((e) => logger.error(e, "Error creating index for assistants by last24HoursCount"));
		assistants
			.createIndex({ last24HoursUseCount: -1, useCount: -1, _id: 1 })
			.catch((e) =>
				logger.error(e, "Error creating index for assistants by last24HoursUseCount and useCount")
			);
		reports
			.createIndex({ assistantId: 1 })
			.catch((e) => logger.error(e, "Error creating index for reports by assistantId"));
		reports
			.createIndex({ createdBy: 1, assistantId: 1 })
			.catch((e) =>
				logger.error(e, "Error creating index for reports by createdBy and assistantId")
			);

		// Unique index for semaphore and migration results
		semaphores.createIndex({ key: 1 }, { unique: true }).catch((e) => logger.error(e));
		semaphores
			.createIndex({ deleteAt: 1 }, { expireAfterSeconds: 1 })
			.catch((e) => logger.error(e, "Error creating index for semaphores by deleteAt"));
		tokenCaches
			.createIndex({ createdAt: 1 }, { expireAfterSeconds: 5 * 60 })
			.catch((e) => logger.error(e, "Error creating index for tokenCaches by createdAt"));
		tokenCaches
			.createIndex({ tokenHash: 1 })
			.catch((e) => logger.error(e, "Error creating index for tokenCaches by tokenHash"));
		conversations
			.createIndex({
				"messages.from": 1,
				createdAt: 1,
			})
			.catch((e) =>
				logger.error(e, "Error creating index for conversations by messages from and createdAt")
			);

		conversations
			.createIndex({
				userId: 1,
				sessionId: 1,
			})
			.catch((e) =>
				logger.error(e, "Error creating index for conversations by userId and sessionId")
			);

		// For stats aggregation jobs that filter by createdAt/updatedAt alone
		conversations
			.createIndex({ createdAt: 1 })
			.catch((e) => logger.error(e, "Error creating index for conversations by createdAt"));
		conversations
			.createIndex({ updatedAt: 1 })
			.catch((e) => logger.error(e, "Error creating index for conversations by updatedAt"));

		config
			.createIndex({ key: 1 }, { unique: true })
			.catch((e) => logger.error(e, "Error creating index for config by key"));
	}
}

export let collections: ReturnType<typeof Database.prototype.getCollections>;

export const ready = (async () => {
	if (!building) {
		const db = await Database.getInstance();
		collections = db.getCollections();
	} else {
		collections = {} as unknown as ReturnType<typeof Database.prototype.getCollections>;
	}
})();

export async function getCollectionsEarly(): Promise<
	ReturnType<typeof Database.prototype.getCollections>
> {
	await ready;
	if (!collections) {
		throw new Error("Database not initialized");
	}
	return collections;
}