@vellumai/assistant 0.10.0 → 0.10.1-dev.202606240317.ea25efe

package/src/memory/migrations/__tests__/run-migrations.test.ts

@@ -3,8 +3,13 @@ import { describe, expect, test } from "bun:test";
 
 import { drizzle } from "drizzle-orm/bun-sqlite";
 
+import { getSqliteFrom } from "../../db-connection.js";
 import * as schema from "../../schema.js";
-import { type MigrationStep, runMigrationSteps } from "../run-migrations.js";
+import {
+  clearMigrationStepCheckpoints,
+  type MigrationStep,
+  runMigrationSteps,
+} from "../run-migrations.js";
 
 function createTestDb() {
   const sqlite = new Database(":memory:");
@@ -14,10 +19,7 @@ function createTestDb() {
 }
 
 describe("runMigrationSteps — checkpointing", () => {
-  // Skipped: documents the post-checkpointing contract and fails against the
-  // current behavior, which re-runs every forward step on every boot. Promote
-  // to `test` when true checkpointing of initial migrations lands.
-  test.skip("does not re-execute a step that was already applied on a prior run", () => {
+  test("does not re-execute a step that was already applied on a prior run", async () => {
     /**
      * Booting an already-migrated database must not re-run forward migration
      * steps. With true checkpointing each step's body executes at most once
@@ -41,12 +43,336 @@ describe("runMigrationSteps — checkpointing", () => {
 
     // AND a database that has already been migrated once
     const db = createTestDb();
-    runMigrationSteps(db, steps);
+    await runMigrationSteps(db, steps);
 
     // WHEN the same steps run again against the already-migrated database
-    runMigrationSteps(db, steps);
+    await runMigrationSteps(db, steps);
 
     // THEN no step body executes a second time
     expect(calls).toEqual({ a: 1, b: 1, c: 1 });
   });
+
+  test("records step completions in the shared memory_checkpoints ledger", async () => {
+    /**
+     * Step bookkeeping lives in the same memory_checkpoints table the registry
+     * uses, under the `step:` namespace — one ledger for all applied state.
+     */
+
+    // GIVEN a single named step
+    const db = createTestDb();
+    await runMigrationSteps(db, [
+      function dummyMigrationA() {
+        // no-op
+      },
+    ]);
+
+    // THEN its completion is recorded under the step: namespace in memory_checkpoints
+    const row = getSqliteFrom(db)
+      .query(
+        `SELECT value FROM memory_checkpoints WHERE key = 'step:dummyMigrationA'`,
+      )
+      .get() as { value: string } | null;
+    expect(row?.value).toBe("1");
+  });
+
+  test("runs a newly appended step on a later boot while skipping applied ones", async () => {
+    /**
+     * Checkpointing must not block migrations added in a later release: an
+     * already-applied step is skipped, but a brand-new step still runs.
+     */
+
+    // GIVEN a database already migrated with two steps
+    const calls = { a: 0, b: 0, c: 0 };
+    const stepA: MigrationStep = function dummyMigrationA() {
+      calls.a++;
+    };
+    const stepB: MigrationStep = function dummyMigrationB() {
+      calls.b++;
+    };
+    const db = createTestDb();
+    await runMigrationSteps(db, [stepA, stepB]);
+
+    // AND a third step appended to the list
+    const stepC: MigrationStep = function dummyMigrationC() {
+      calls.c++;
+    };
+
+    // WHEN the expanded list runs against the already-migrated database
+    await runMigrationSteps(db, [stepA, stepB, stepC]);
+
+    // THEN only the newly appended step executes
+    expect(calls).toEqual({ a: 1, b: 1, c: 1 });
+  });
+
+  test("recovers crashed migrations before running steps", async () => {
+    /**
+     * runMigrationSteps clears stalled `started`/`rolling_back` checkpoints
+     * before the loop so a migration interrupted mid-flight re-runs this boot.
+     */
+
+    // GIVEN a database with a stalled registry checkpoint left by a crash
+    const db = createTestDb();
+    const raw = getSqliteFrom(db);
+    raw.run(
+      `CREATE TABLE IF NOT EXISTS memory_checkpoints (key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at INTEGER NOT NULL)`,
+    );
+    raw.run(
+      `INSERT INTO memory_checkpoints (key, value, updated_at) VALUES ('migration_stalled_v1', 'started', 0)`,
+    );
+
+    // WHEN the runner executes
+    await runMigrationSteps(db, [
+      function dummyMigrationA() {
+        // no-op
+      },
+    ]);
+
+    // THEN the stalled checkpoint has been cleared so its migration can re-run
+    const row = raw
+      .query(
+        `SELECT 1 FROM memory_checkpoints WHERE key = 'migration_stalled_v1'`,
+      )
+      .get();
+    expect(row).toBeNull();
+  });
+
+  test("does not checkpoint a failed step so it retries on the next boot", async () => {
+    /**
+     * A step whose body throws is reported in `failed` and left uncheckpointed,
+     * so the next boot retries it instead of silently skipping it.
+     */
+
+    // GIVEN a step that throws on its first run and succeeds afterwards
+    const calls = { flaky: 0 };
+    const steps: MigrationStep[] = [
+      function flakyStep() {
+        calls.flaky++;
+        if (calls.flaky === 1) {
+          throw new Error("transient failure");
+        }
+      },
+    ];
+    const db = createTestDb();
+    const first = await runMigrationSteps(db, steps);
+
+    // AND the first run reports the failure
+    expect(first.failed).toEqual(["flakyStep"]);
+
+    // WHEN the step runs again on the next boot
+    const second = await runMigrationSteps(db, steps);
+
+    // THEN it retries, succeeds, and is then checkpointed
+    expect(calls.flaky).toBe(2);
+    expect(second.failed).toEqual([]);
+    expect((await runMigrationSteps(db, steps)).skipped).toEqual(["flakyStep"]);
+  });
+
+  test("clearMigrationStepCheckpoints forces every step to re-run", async () => {
+    /**
+     * Clearing the step namespace (as rollback does) makes the next run
+     * re-execute and re-record all steps.
+     */
+
+    // GIVEN a database migrated once
+    const calls = { a: 0 };
+    const steps: MigrationStep[] = [
+      function dummyMigrationA() {
+        calls.a++;
+      },
+    ];
+    const db = createTestDb();
+    await runMigrationSteps(db, steps);
+    expect(calls.a).toBe(1);
+
+    // WHEN the step checkpoints are cleared
+    clearMigrationStepCheckpoints(db);
+
+    // THEN the next run re-executes the step
+    await runMigrationSteps(db, steps);
+    expect(calls.a).toBe(2);
+  });
+
+  test("awaits an async step to completion before running or checkpointing the next", async () => {
+    /**
+     * An async step must be fully drained before the runner advances: ordering
+     * is the invariant later migrations rely on (step N+1 may assume N's result
+     * is true), and the step is checkpointed only after its promise resolves so
+     * a crash mid-drain leaves it uncheckpointed and retryable rather than
+     * recorded as done.
+     */
+
+    // GIVEN an async step that records its completion order only after yielding,
+    // followed by a sync step that records when it starts
+    const order: string[] = [];
+    const steps: MigrationStep[] = [
+      async function asyncDrainStep() {
+        await Promise.resolve();
+        order.push("async-done");
+      },
+      function followingStep() {
+        order.push("following-start");
+      },
+    ];
+    const db = createTestDb();
+
+    // WHEN the runner executes both steps
+    const result = await runMigrationSteps(db, steps);
+
+    // THEN the async step finishes before the next step starts
+    expect(order).toEqual(["async-done", "following-start"]);
+
+    // AND both steps are checkpointed, so a later boot skips them
+    expect(result.failed).toEqual([]);
+    expect((await runMigrationSteps(db, steps)).skipped).toEqual([
+      "asyncDrainStep",
+      "followingStep",
+    ]);
+  });
+
+  test("a rejected async step is reported failed and left uncheckpointed", async () => {
+    /**
+     * A step whose promise rejects must be treated exactly like a synchronous
+     * throw: reported in `failed`, not checkpointed, and retried next boot.
+     */
+
+    // GIVEN an async step that rejects on its first run and resolves afterwards
+    const calls = { flaky: 0 };
+    const steps: MigrationStep[] = [
+      async function flakyAsyncStep() {
+        calls.flaky++;
+        await Promise.resolve();
+        if (calls.flaky === 1) {
+          throw new Error("transient async failure");
+        }
+      },
+    ];
+    const db = createTestDb();
+
+    // WHEN it runs and rejects
+    const first = await runMigrationSteps(db, steps);
+
+    // THEN the rejection is reported and the step is not checkpointed
+    expect(first.failed).toEqual(["flakyAsyncStep"]);
+
+    // WHEN it runs again on the next boot
+    const second = await runMigrationSteps(db, steps);
+
+    // THEN it retries, resolves, and is then checkpointed
+    expect(calls.flaky).toBe(2);
+    expect(second.failed).toEqual([]);
+    expect((await runMigrationSteps(db, steps)).skipped).toEqual([
+      "flakyAsyncStep",
+    ]);
+  });
+
+  test("detects and recovers step-level 'started' markers from a prior crash", async () => {
+    /**
+     * If the process crashes mid-step, the runner leaves a `step:<name>` =
+     * 'started' checkpoint. On the next boot, recoverCrashedMigrations
+     * detects it, clears it, and the step re-runs.
+     */
+
+    // GIVEN a database with a stalled step: checkpoint left by a crash
+    const db = createTestDb();
+    const raw = getSqliteFrom(db);
+    raw.run(
+      `CREATE TABLE IF NOT EXISTS memory_checkpoints (key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at INTEGER NOT NULL)`,
+    );
+    raw.run(
+      `INSERT INTO memory_checkpoints (key, value, updated_at) VALUES ('step:crashedStep', 'started', 0)`,
+    );
+
+    let ran = false;
+    const steps: MigrationStep[] = [
+      function crashedStep() {
+        ran = true;
+      },
+    ];
+
+    // WHEN the runner executes — recoverCrashedMigrations should clear
+    // the 'started' marker and let the step run.
+    await runMigrationSteps(db, steps);
+
+    // THEN the stalled checkpoint was cleared and the step ran
+    expect(ran).toBe(true);
+    const stalled = raw
+      .query(
+        `SELECT 1 FROM memory_checkpoints WHERE key = 'step:crashedStep' AND value = 'started'`,
+      )
+      .get();
+    expect(stalled).toBeNull();
+    // And the step was checkpointed as '1' (completed)
+    const completed = raw
+      .query(
+        `SELECT 1 FROM memory_checkpoints WHERE key = 'step:crashedStep' AND value = '1'`,
+      )
+      .get();
+    expect(completed).not.toBeNull();
+  });
+
+  test("does not skip a step whose only checkpoint is 'started' (not '1')", async () => {
+    /**
+     * The skip query filters for value = '1', so a 'started' marker
+     * alone must not cause a step to be skipped — the step should run.
+     * (recoverCrashedMigrations clears 'started' markers first, but
+     * even if it didn't, the skip logic is independent and correct.)
+     */
+
+    const db = createTestDb();
+    const raw = getSqliteFrom(db);
+    raw.run(
+      `CREATE TABLE IF NOT EXISTS memory_checkpoints (key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at INTEGER NOT NULL)`,
+    );
+    // Insert a 'started' marker without running recoverCrashedMigrations.
+    // We simulate this by inserting directly and checking the applied set.
+    raw.run(
+      `INSERT INTO memory_checkpoints (key, value, updated_at) VALUES ('step:someStep', 'started', 0)`,
+    );
+
+    let ran = false;
+    const steps: MigrationStep[] = [
+      function someStep() {
+        ran = true;
+      },
+    ];
+
+    // WHEN the runner executes, recoverCrashedMigrations clears the
+    // 'started' marker, so the step is not in the applied set and runs.
+    const result = await runMigrationSteps(db, steps);
+
+    // THEN the step ran (was not skipped)
+    expect(ran).toBe(true);
+    expect(result.applied).toEqual(["someStep"]);
+    expect(result.skipped).toEqual([]);
+  });
+
+  test("writes 'started' marker before running each step", async () => {
+    const db = createTestDb();
+    const raw = getSqliteFrom(db);
+
+    const seen: { marker: string | null } = { marker: null };
+    const steps: MigrationStep[] = [
+      async function asyncWithInspection() {
+        const row = raw
+          .query<{ value: string }, []>(
+            `SELECT value FROM memory_checkpoints WHERE key = 'step:asyncWithInspection'`,
+          )
+          .get();
+        seen.marker = row?.value ?? null;
+        await Promise.resolve();
+      },
+    ];
+
+    await runMigrationSteps(db, steps);
+
+    // The 'started' marker was visible during the async step's execution
+    expect(seen.marker).toBe("started");
+
+    const final = raw
+      .query<{ value: string }, []>(
+        `SELECT value FROM memory_checkpoints WHERE key = 'step:asyncWithInspection'`,
+      )
+      .get();
+    expect(final?.value).toBe("1");
+  });
 });

package/src/memory/migrations/helpers/relocation.ts

@@ -0,0 +1,227 @@
+import type { Database } from "bun:sqlite";
+
+import { getLogger } from "../../../util/logger.js";
+import { getDbPath } from "../../../util/platform.js";
+import {
+  parseChangesFromStdout,
+  runAsyncSqlite,
+} from "../../db-async-query.js";
+
+const log = getLogger("memory-db");
+
+/**
+ * Incremental relocation of a heavy table out of the main DB and into its
+ * own file (`assistant-logs.db` / `assistant-memory.db`).
+ *
+ * A one-shot `INSERT … SELECT` + `DROP TABLE` of a multi-GB table on the daemon
+ * connection would pin the write lock and block the event loop for minutes.
+ * Instead a relocation runs as an async migration step in two parts:
+ *
+ *   1. {@link stageTableForRelocation} renames the source table in `main` aside
+ *      to `<table>__relocating` (instant, metadata-only), so live reads/writes
+ *      route to the dedicated connection's copy immediately.
+ *   2. {@link drainStagedTable} copies the staged rows into the target file in
+ *      bounded batches and truncates them from the staging table as it goes,
+ *      then drops it. Each batch runs through `runAsyncSqlite`, which opens the
+ *      target file directly (the sqlite3 subprocess ATTACHes it), so the work
+ *      is independent of the daemon connection; the migration `await`s it to
+ *      completion before checkpointing, so later startup work observes the
+ *      finished move.
+ *
+ * The engine is generic: each migration owns its {@link RelocationSpec} (the
+ * instance-specific columns / filter / target file) and passes it in.
+ */
+
+/** Suffix of the staging table the source is renamed to during relocation. */
+export const RELOCATING_SUFFIX = "__relocating";
+
+/**
+ * Describes how to drain one relocatable table. Defined by the owning migration
+ * — never centrally — so instance-specific SQL stays next to the migration that
+ * needs it.
+ */
+export interface RelocationSpec {
+  /** Live, unqualified table name (renamed aside, then drained into the target). */
+  table: string;
+  /** Absolute path of the attached DB file the table is moved into. */
+  targetDbPath: () => string;
+  /**
+   * Columns to copy, in a fixed order — listed explicitly (not `SELECT *`) so
+   * the copy is insensitive to the physical column order of the source table,
+   * which varies with the history of `ALTER TABLE … ADD COLUMN` migrations.
+   * Columns absent from the (legacy) source are copied as NULL.
+   */
+  columns: string[];
+  /**
+   * Optional predicate (evaluated against the staging table) selecting rows
+   * worth preserving. Rows that do **not** match are deleted without being
+   * copied. Omit to copy every row.
+   */
+  copyWhere?: string;
+  /**
+   * Optional per-column SELECT expressions, keyed by column name, for columns
+   * that must be transformed during the copy rather than carried verbatim
+   * (e.g. resetting a status). The expression is evaluated against the staging
+   * row; unlisted columns copy as-is (or NULL when absent from the source).
+   */
+  columnExpr?: Record<string, string>;
+}
+
+function tableExistsInMain(raw: Database, name: string): boolean {
+  return (
+    raw
+      .query(
+        `SELECT name FROM main.sqlite_master WHERE type='table' AND name = ?`,
+      )
+      .get(name) != null
+  );
+}
+
+function tableIsEmpty(raw: Database, name: string): boolean {
+  // EXISTS short-circuits at the first row, so this stays cheap even on a huge
+  // table — no full COUNT(*) scan.
+  return raw.query(`SELECT 1 FROM main."${name}" LIMIT 1`).get() == null;
+}
+
+/**
+ * Move the source table in `main` aside so the unqualified name resolves to the
+ * attached copy, returning whether a staging table now exists (i.e. a drain is
+ * needed). Idempotent and safe to re-run after a crash:
+ *
+ *   - `main.<table>` empty           → drop it (a freshly recreated shadow, or
+ *                                       an already-drained leftover).
+ *   - `main.<table>` non-empty, no
+ *     staging yet                     → rename it to `<table>__relocating`.
+ *   - staging already exists          → leave it; a prior boot started the move.
+ *
+ * `table` comes from a {@link RelocationSpec} (never user input); it is quoted
+ * defensively all the same.
+ */
+export function stageTableForRelocation(raw: Database, table: string): boolean {
+  const staging = `${table}${RELOCATING_SUFFIX}`;
+  const hasStaging = tableExistsInMain(raw, staging);
+
+  if (tableExistsInMain(raw, table)) {
+    if (tableIsEmpty(raw, table)) {
+      raw.exec(`DROP TABLE main."${table}"`);
+    } else if (!hasStaging) {
+      raw.exec(`ALTER TABLE main."${table}" RENAME TO "${staging}"`);
+      return true;
+    }
+    // else: a non-empty live table alongside an existing staging table is not
+    // expected; leave both and let the in-flight drain finish first.
+  }
+
+  return tableExistsInMain(raw, staging);
+}
+
+/**
+ * Rows copied/purged per drain batch. Each batch is a couple of bounded
+ * statements, so the write lock is held only briefly and the event loop is free
+ * between batches. Sized as a balance between throughput and lock-hold time.
+ */
+const DRAIN_BATCH = 10_000;
+
+/**
+ * Drain a `<table>__relocating` staging table created by
+ * {@link stageTableForRelocation}: copy the rows worth keeping into the attached
+ * target in bounded batches, purge the rest without copying, then drop the
+ * staging table and truncate the main WAL. Resolves once the move is complete.
+ *
+ * Awaited inline by the relocation migration step, so it runs at most once per
+ * boot under the migration runner's checkpoint: an interrupted drain leaves the
+ * step uncheckpointed and the staging table in place, and the next boot
+ * re-stages (a no-op) and resumes from the remaining rows. Each batch is
+ * dispatched off the connection via `runAsyncSqlite` (sqlite3 subprocess when
+ * available; in-process fallback otherwise), keeping the event loop responsive
+ * between batches. A batch failure throws so the step is reported failed and
+ * retried on the next boot rather than checkpointed as done.
+ */
+export async function drainStagedTable(
+  raw: Database,
+  spec: RelocationSpec,
+): Promise<void> {
+  const { table } = spec;
+  const staging = `${table}${RELOCATING_SUFFIX}`;
+
+  // Nothing to do once the staging table is gone (drain finished previously).
+  if (!tableExistsInMain(raw, staging)) return;
+
+  // Build a select list from the staging table's actual columns: apply any
+  // per-column transform, copy a present column verbatim, NULL-fill an absent
+  // (legacy) one so an older row still copies cleanly.
+  const present = new Set(
+    (
+      raw
+        .query(`SELECT name FROM pragma_table_info('${staging}', 'main')`)
+        .all() as Array<{ name: string }>
+    ).map((r) => r.name),
+  );
+  const colList = spec.columns.map((c) => `"${c}"`).join(", ");
+  const selectList = spec.columns
+    .map((c) => spec.columnExpr?.[c] ?? (present.has(c) ? `"${c}"` : "NULL"))
+    .join(", ");
+
+  const dbPath = getDbPath();
+  const whereCopy = spec.copyWhere ? `WHERE ${spec.copyWhere}` : "";
+
+  for (;;) {
+    // (1) Purge a batch of non-keeper rows (no copy) when a copy filter narrows
+    //     what is worth preserving — this is the bulk of a runaway queue.
+    let purged = 0;
+    if (spec.copyWhere) {
+      const res = await runAsyncSqlite(
+        `DELETE FROM "${staging}" WHERE rowid IN (` +
+          `SELECT rowid FROM "${staging}" WHERE NOT (${spec.copyWhere}) LIMIT ${DRAIN_BATCH});\n` +
+          `SELECT changes();`,
+        { dbPath },
+      );
+      if (!res.ok) {
+        throw new Error(
+          `relocation purge batch failed for "${table}": ${res.error}`,
+        );
+      }
+      purged = parseChangesFromStdout(res.stdout);
+    }
+
+    // (2) Copy a batch of keepers, then truncate those same rows. Two
+    //     autocommitted statements (no BEGIN): the copy commits before the
+    //     delete, so a crash in between just re-copies (INSERT OR IGNORE no-op)
+    //     and re-deletes next boot — safe across the non-atomic cross-DB commit.
+    const copyRes = await runAsyncSqlite(
+      `INSERT OR IGNORE INTO "${table}" (${colList}) ` +
+        `SELECT ${selectList} FROM "${staging}" ${whereCopy} ORDER BY rowid LIMIT ${DRAIN_BATCH};\n` +
+        `DELETE FROM "${staging}" WHERE rowid IN (` +
+        `SELECT rowid FROM "${staging}" ${whereCopy} ORDER BY rowid LIMIT ${DRAIN_BATCH});\n` +
+        `SELECT changes();`,
+      {
+        dbPath,
+        attach: [{ path: spec.targetDbPath(), alias: "_reloc_target" }],
+      },
+    );
+    if (!copyRes.ok) {
+      throw new Error(
+        `relocation copy batch failed for "${table}": ${copyRes.error}`,
+      );
+    }
+    const moved = parseChangesFromStdout(copyRes.stdout);
+
+    if (purged > 0 || moved > 0) {
+      log.info({ table, purged, moved }, "relocation: drain progressed");
+      continue;
+    }
+    break;
+  }
+
+  // Drained — drop the (now empty) staging table and truncate the main WAL.
+  const finalizeRes = await runAsyncSqlite(
+    `DROP TABLE IF EXISTS "${staging}";\nPRAGMA wal_checkpoint(TRUNCATE);`,
+    { dbPath },
+  );
+  if (!finalizeRes.ok) {
+    throw new Error(
+      `relocation finalize failed for "${table}": ${finalizeRes.error}`,
+    );
+  }
+  log.info({ table }, "relocation: complete — staging dropped");
+}