Add bots files handling logic

src/file-ingest/enqueue-file-ingest.ts

@@ -0,0 +1,37 @@
+import BunSQLite from "@moduletrace/bun-sqlite";
+import type { BUN_SQLITE_COBBLE_FILE_INGEST_QUEUE } from "@/db/types/db";
+import wakeFileIngest from "./wake-file-ingest";
+
+type Params = {
+    file_id: number | string;
+};
+
+/**
+ * Queues a file to be chunked and embedded, unless it's already waiting.
+ * Never fails the caller: a file that can't be queued is picked up by the
+ * `ingest-pending-files` cron instead.
+ */
+export default async function enqueueFileIngest({ file_id }: Params) {
+    try {
+        const open = await BunSQLite.sql<BUN_SQLITE_COBBLE_FILE_INGEST_QUEUE>({
+            sql: `SELECT id FROM file_ingest_queue WHERE file_id = ? AND status IN ('pending', 'processing') LIMIT 1`,
+            values: [file_id],
+        });
+
+        if (open.payload?.[0]?.id) {
+            wakeFileIngest();
+            return;
+        }
+
+        await BunSQLite.queue.add({
+            table: "file_ingest_queue",
+            data: [{ file_id: Number(file_id) }],
+        });
+
+        wakeFileIngest();
+    } catch {
+        /**
+         * The cron sweeps it up.
+         */
+    }
+}

src/file-ingest/index.ts

@@ -0,0 +1,82 @@
+import bunext from "@moduletrace/bunext";
+import BunSQLite from "@moduletrace/bun-sqlite";
+import type { BUN_SQLITE_COBBLE_FILE_INGEST_QUEUE } from "@/db/types/db";
+import { AppData } from "@/src/data/app-data";
+import { logError } from "@/src/functions/backend/logger";
+import processIngestJob from "./process-ingest-job";
+import { setFileIngestWaker } from "./wake-file-ingest";
+
+const worker = AppData.FileIngestWorker;
+let running = 0;
+let polling = false;
+let timer: ReturnType<typeof setTimeout> | undefined;
+
+/**
+ * Starts the file ingestion worker: polls every `AppData.FileIngestPollMs`
+ * for `file_ingest_queue` jobs and runs up to
+ * `AppData.FileIngestConcurrency` at once (chunking and embedding a bot's
+ * file). Gated in production to the `MAIN_SERVER` container, like cron and
+ * the bots queue.
+ */
+export default async function fileIngest() {
+    if (process.env.NODE_ENV == "production" && !process.env.MAIN_SERVER) {
+        return;
+    }
+
+    bunext.bunextLog.watch(`Running File Ingest Queue ...`);
+
+    /**
+     * A restart's cut-off jobs go back to waiting (their lock timeout
+     * would, but this is immediate).
+     */
+    await BunSQLite.queue.release_stale({ table: "file_ingest_queue" }).catch((error) =>
+        logError({ title: "File ingest couldn't release stale jobs", error }).catch(() => {}),
+    );
+
+    setFileIngestWaker({ waker: wake });
+    poll();
+}
+
+function wake() {
+    if (polling) return;
+
+    clearTimeout(timer);
+    poll();
+}
+
+async function poll() {
+    polling = true;
+
+    try {
+        const free = AppData.FileIngestConcurrency - running;
+
+        if (free > 0) {
+            const { payload } = await BunSQLite.queue.claim<BUN_SQLITE_COBBLE_FILE_INGEST_QUEUE>({
+                table: "file_ingest_queue",
+                worker,
+                limit: free,
+            });
+
+            for (const job of payload || []) {
+                running++;
+
+                processIngestJob({ job })
+                    .catch((error) =>
+                        logError({
+                            title: `File ingest job #${job.id} crashed`,
+                            error,
+                        }).catch(() => {}),
+                    )
+                    .finally(() => {
+                        running--;
+                    });
+            }
+        }
+    } catch (error) {
+        await logError({ title: "File ingest poll failed", error }).catch(() => {});
+    } finally {
+        polling = false;
+        clearTimeout(timer);
+        timer = setTimeout(poll, AppData.FileIngestPollMs);
+    }
+}

src/file-ingest/process-ingest-job.ts

@@ -0,0 +1,41 @@
+import BunSQLite from "@moduletrace/bun-sqlite";
+import type { BUN_SQLITE_COBBLE_FILE_INGEST_QUEUE } from "@/db/types/db";
+import ingestBotFile from "@/src/functions/backend/db/bot-files/ingest-bot-file";
+import { logError } from "@/src/functions/backend/logger";
+
+type Params = {
+    job: BUN_SQLITE_COBBLE_FILE_INGEST_QUEUE;
+};
+
+/**
+ * Ingests one queued file and completes (or fails) its job. A file that
+ * was deleted meanwhile is simply done.
+ */
+export default async function processIngestJob({ job }: Params) {
+    if (!job.id || !job.file_id) return;
+
+    try {
+        const result = await ingestBotFile({ file_id: job.file_id });
+
+        await BunSQLite.queue.complete({ table: "file_ingest_queue", id: job.id });
+
+        if (result.status == "ingested") {
+            BunSQLite.sql({
+                sql: `UPDATE bot_files SET ingest_error = NULL WHERE id = ?`,
+                values: [job.file_id],
+            }).catch(() => {});
+        }
+    } catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+
+        await BunSQLite.queue
+            .fail({ table: "file_ingest_queue", id: job.id, error: message.slice(0, 500) })
+            .catch(() => {});
+
+        await logError({
+            title: `Bot file #${job.file_id} couldn't be ingested`,
+            error,
+            metadata: { file_id: job.file_id },
+        }).catch(() => {});
+    }
+}

src/file-ingest/wake-file-ingest.ts

@@ -0,0 +1,17 @@
+let wake: (() => void) | undefined;
+
+/**
+ * Lets the running ingest worker register how to poll at once, kept apart
+ * from the handler so an upload can wake it without bundling it.
+ */
+export function setFileIngestWaker({ waker }: { waker: () => void }) {
+    wake = waker;
+}
+
+/**
+ * Polls now instead of at the next tick, so a fresh file is readable
+ * moments after it lands. Does nothing where the worker doesn't run.
+ */
+export default function wakeFileIngest() {
+    wake?.();
+}