@vellumai/assistant 0.10.0 → 0.10.1-staging.1

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

@@ -1,41 +1,212 @@
 import { getLogger } from "../../util/logger.js";
-import type { DrizzleDb } from "../db-connection.js";
+import { getDbPath } from "../../util/platform.js";
+import { type DrizzleDb, getSqliteFrom } from "../db-connection.js";
 
 const log = getLogger("db-init");
 
-/** A single forward migration step, identified for logging by its `.name`. */
-export type MigrationStep = (database: DrizzleDb) => void;
+/**
+ * A single forward migration step, identified for checkpointing and logging by
+ * its `.name`. Anonymous steps (empty `.name`) cannot be tracked and always run.
+ *
+ * A step may be synchronous or return a promise. The runner awaits an async
+ * step to completion before checkpointing it and moving on, so ordering is
+ * preserved exactly as for sync steps: step N+1 never starts — and is never
+ * skipped via checkpoint — until step N has fully finished. This lets a step
+ * that drains a large backfill in `await`ed batches run without blocking the
+ * thread between batches while still guaranteeing later migrations observe its
+ * completed result.
+ */
+export type MigrationStep = (database: DrizzleDb) => void | Promise<void>;
 
 export interface MigrationRunResult {
+  /** Steps that ran and completed successfully this boot. */
+  applied: string[];
   /** Steps whose body threw. */
   failed: string[];
+  /** Steps skipped because a prior run already applied them. */
+  skipped: string[];
+}
+
+/**
+ * Prefix under which forward-step completions are recorded in the shared
+ * `memory_checkpoints` ledger. A distinct namespace keeps step bookkeeping from
+ * colliding with registry checkpoint keys (`migration_*`, `backfill_*`,
+ * `drop_*`) and is deliberately chosen so `validateMigrationState` does not
+ * mistake a step record for an unknown registry migration.
+ */
+export const STEP_CHECKPOINT_PREFIX = "step:";
+
+/**
+ * Create the migration bookkeeping table if it is missing.
+ *
+ * The runner records its own checkpoints, so it must not depend on an earlier
+ * step having created the ledger. This is the single place `memory_checkpoints`
+ * is created; `IF NOT EXISTS` makes it a no-op on an already-migrated database.
+ */
+function ensureCheckpointsTable(raw: ReturnType<typeof getSqliteFrom>): void {
+  raw.run(/*sql*/ `
+    CREATE TABLE IF NOT EXISTS memory_checkpoints (
+      key TEXT PRIMARY KEY,
+      value TEXT NOT NULL,
+      updated_at INTEGER NOT NULL
+    )
+  `);
+}
+
+/**
+ * Recover from crashed migrations before the migration runner executes its steps.
+ *
+ * Scans memory_checkpoints for entries with value 'started' or 'rolling_back' —
+ * these represent migrations that began but never completed (e.g., due to a
+ * process crash). Deletes the stalled checkpoint so the migration can re-run
+ * from scratch on this startup. Each migration's own idempotency guards (DDL
+ * IF NOT EXISTS, transactional rollback) ensure re-running is safe.
+ *
+ * Runs on every boot — it must observe the state left by *this* boot's prior
+ * crash — so it is invoked directly by {@link runMigrationSteps} before the
+ * checkpointed step loop rather than being a checkpointed step itself.
+ */
+export function recoverCrashedMigrations(database: DrizzleDb): string[] {
+  const raw = getSqliteFrom(database);
+
+  let rows: Array<{ key: string; value: string }>;
+  try {
+    rows = raw
+      .query(`SELECT key, value FROM memory_checkpoints`)
+      .all() as Array<{ key: string; value: string }>;
+  } catch {
+    return [];
+  }
+
+  const crashed = rows
+    .filter((r) => r.value === "started" || r.value === "rolling_back")
+    .map((r) => r.key);
+  if (crashed.length === 0) return [];
+
+  log.error(
+    { crashed },
+    [
+      "╔══════════════════════════════════════════════════════════════╗",
+      "║  CRASHED MIGRATIONS DETECTED — AUTO-RECOVERING             ║",
+      "╚══════════════════════════════════════════════════════════════╝",
+      "",
+      `The following migrations started but never completed: ${crashed.join(", ")}`,
+      "",
+      "Clearing stalled checkpoints so they can be retried on this startup.",
+      "If retries continue to fail, manually inspect the database:",
+      `  sqlite3 ${getDbPath()} "SELECT * FROM memory_checkpoints"`,
+    ].join("\n"),
+  );
+
+  for (const key of crashed) {
+    raw.query(`DELETE FROM memory_checkpoints WHERE key = ?`).run(key);
+    log.info(
+      { key },
+      `Cleared stalled checkpoint "${key}" — migration will re-run`,
+    );
+  }
+
+  return crashed;
 }
 
 /**
- * Run the ordered list of forward migration steps against the database.
+ * Run the ordered list of forward migration steps against the database, each at
+ * most once across boots.
+ *
+ * A step's name is recorded in `memory_checkpoints` after its body completes —
+ * for an async step, after its returned promise resolves; on
+ * later boots an applied step is skipped instead of re-executed. This turns the
+ * unconditional ~200-step re-probe — which floors daemon startup at tens of
+ * seconds on a fully-migrated database — into a single bookkeeping read.
+ *
+ * Step bookkeeping lives in the `memory_checkpoints` ledger under the
+ * {@link STEP_CHECKPOINT_PREFIX} namespace, so applied-state for every
+ * migration lives in one place.
+ *
+ * Before the step loop runs, the ledger is created if missing and
+ * {@link recoverCrashedMigrations} clears any stalled checkpoints left by a
+ * prior crash, so a migration interrupted mid-flight re-runs this boot.
  *
  * Individual step failures are caught and logged so one broken migration does
- * not prevent independent later ones from succeeding.
+ * not prevent independent later ones from succeeding; a failed step is not
+ * checkpointed and is retried on the next boot.
  */
-export function runMigrationSteps(
+export async function runMigrationSteps(
   database: DrizzleDb,
   steps: MigrationStep[],
-): MigrationRunResult {
+): Promise<MigrationRunResult> {
+  const raw = getSqliteFrom(database);
+
+  ensureCheckpointsTable(raw);
+  recoverCrashedMigrations(database);
+
+  const applied = new Set(
+    (
+      raw
+        .query(
+          `SELECT key FROM memory_checkpoints WHERE key LIKE '${STEP_CHECKPOINT_PREFIX}%' AND value = '1'`,
+        )
+        .all() as Array<{ key: string }>
+    ).map((row) => row.key.slice(STEP_CHECKPOINT_PREFIX.length)),
+  );
+  const markStarted = raw.query(
+    `INSERT OR REPLACE INTO memory_checkpoints (key, value, updated_at) VALUES (?, 'started', ?)`,
+  );
+  const markApplied = raw.query(
+    `INSERT OR REPLACE INTO memory_checkpoints (key, value, updated_at) VALUES (?, '1', ?)`,
+  );
+
   const failed: string[] = [];
+  const skipped: string[] = [];
+  const ran: string[] = [];
 
   for (const step of steps) {
+    const name = step.name;
+    const checkpointable = name !== "";
+
+    if (checkpointable && applied.has(name)) {
+      skipped.push(name);
+      log.debug({ migration: name }, `Skipping applied migration: ${name}`);
+      continue;
+    }
+
     try {
-      log.debug({ migration: step.name }, `Starting migration: ${step.name}`);
-      step(database);
-      log.debug({ migration: step.name }, `Migration succeeded: ${step.name}`);
+      if (checkpointable) {
+        markStarted.run(`${STEP_CHECKPOINT_PREFIX}${name}`, Date.now());
+      }
+      log.info({ migration: name }, `Starting migration: ${name}`);
+      const result = step(database);
+      if (result instanceof Promise) {
+        await result;
+      }
+      log.info({ migration: name }, `Migration succeeded: ${name}`);
+      if (checkpointable) {
+        markApplied.run(`${STEP_CHECKPOINT_PREFIX}${name}`, Date.now());
+        ran.push(name);
+      }
     } catch (err) {
-      failed.push(step.name);
-      log.error(
-        { err, migration: step.name },
-        `Migration failed: ${step.name}`,
-      );
+      // Leave the 'started' marker in place (if one was written) —
+      // recoverCrashedMigrations will detect it on the next boot, log
+      // a warning, and clear it so the step re-runs.
+      failed.push(name);
+      log.error({ err, migration: name }, `Migration failed: ${name}`);
     }
   }
 
-  return { failed };
+  return { applied: ran, failed, skipped };
+}
+
+/**
+ * Discard every forward-step checkpoint so the next {@link runMigrationSteps}
+ * call re-runs and re-records all steps.
+ *
+ * Migration rollback calls this. A rolled-back step's `step:` checkpoint
+ * must be discarded, otherwise the runner skips the step on the next
+ * upgrade and the rolled-back schema is never restored. Only the `step:`
+ * namespace is cleared, leaving other ledger entries untouched.
+ */
+export function clearMigrationStepCheckpoints(database: DrizzleDb): void {
+  getSqliteFrom(database).run(
+    `DELETE FROM memory_checkpoints WHERE key LIKE '${STEP_CHECKPOINT_PREFIX}%'`,
+  );
 }

package/src/memory/migrations/validate-migration-state.ts

@@ -6,120 +6,12 @@ import {
   MIGRATION_REGISTRY,
   type MigrationValidationResult,
 } from "./registry.js";
+import {
+  STEP_CHECKPOINT_PREFIX,
+} from "./run-migrations.js";
 
 const log = getLogger("memory-db");
 
-/**
- * Recover from crashed migrations before the migration runner executes.
- *
- * Scans memory_checkpoints for entries with value 'started' — these represent
- * migrations that began but never completed (e.g., due to a process crash).
- * Deletes the stalled checkpoint so the migration can re-run from scratch on
- * this startup. Each migration's own idempotency guards (DDL IF NOT EXISTS,
- * transactional rollback) ensure re-running is safe.
- *
- * Call this BEFORE running migrations so that stalled checkpoints don't block
- * re-execution.
- */
-export function recoverCrashedMigrations(database: DrizzleDb): string[] {
-  const raw = getSqliteFrom(database);
-
-  let rows: Array<{ key: string; value: string }>;
-  try {
-    rows = raw
-      .query(`SELECT key, value FROM memory_checkpoints`)
-      .all() as Array<{ key: string; value: string }>;
-  } catch {
-    return [];
-  }
-
-  const crashed = rows
-    .filter((r) => r.value === "started" || r.value === "rolling_back")
-    .map((r) => r.key);
-  if (crashed.length === 0) return [];
-
-  log.error(
-    { crashed },
-    [
-      "╔══════════════════════════════════════════════════════════════╗",
-      "║  CRASHED MIGRATIONS DETECTED — AUTO-RECOVERING             ║",
-      "╚══════════════════════════════════════════════════════════════╝",
-      "",
-      `The following migrations started but never completed: ${crashed.join(", ")}`,
-      "",
-      "Clearing stalled checkpoints so they can be retried on this startup.",
-      "If retries continue to fail, manually inspect the database:",
-      `  sqlite3 ${getDbPath()} "SELECT * FROM memory_checkpoints"`,
-    ].join("\n"),
-  );
-
-  for (const key of crashed) {
-    raw.query(`DELETE FROM memory_checkpoints WHERE key = ?`).run(key);
-    log.info(
-      { key },
-      `Cleared stalled checkpoint "${key}" — migration will re-run`,
-    );
-  }
-
-  return crashed;
-}
-
-/**
- * Wrap a migration function with crash-recovery bookkeeping.
- *
- * Writes a 'started' checkpoint before executing the migration body, then
- * overwrites it with the completion value on success. If the process crashes
- * between the start marker and completion, recoverCrashedMigrations (which
- * runs before all migrations) will detect and clear it on the next startup.
- *
- * The migrationFn receives the raw SQLite database and should perform its
- * own transaction management internally.
- */
-export function withCrashRecovery(
-  database: DrizzleDb,
-  checkpointKey: string,
-  migrationFn: () => void,
-): void {
-  const raw = getSqliteFrom(database);
-
-  const existing = raw
-    .query(`SELECT value FROM memory_checkpoints WHERE key = ?`)
-    .get(checkpointKey) as { value: string } | null;
-  if (
-    existing &&
-    existing.value !== "started" &&
-    existing.value !== "rolling_back"
-  )
-    return;
-
-  raw
-    .query(
-      `INSERT OR REPLACE INTO memory_checkpoints (key, value, updated_at) VALUES (?, 'started', ?)`,
-    )
-    .run(checkpointKey, Date.now());
-
-  try {
-    migrationFn();
-  } catch (error) {
-    log.error(
-      { checkpointKey, error },
-      `Memory migration failed: ${checkpointKey} — marking as failed and continuing`,
-    );
-    raw
-      .query(
-        `UPDATE memory_checkpoints SET value = 'failed', updated_at = ? WHERE key = ?`,
-      )
-      .run(Date.now(), checkpointKey);
-    return;
-  }
-
-  raw
-    .query(
-      `UPDATE memory_checkpoints SET value = '1', updated_at = ? WHERE key = ?`,
-    )
-    .run(Date.now(), checkpointKey);
-}
-
 /**
  * Validate the applied migration state against the registry at startup.
  *
@@ -171,16 +63,28 @@ export function validateMigrationState(
     );
   }
 
-  // Only rows whose value is NOT 'started' or 'rolling_back' represent truly
-  // completed migrations. In-progress/crashed checkpoints must not count as
-  // applied dependencies — the migration never finished, so its postconditions
-  // are unmet.
-  const completed = new Set(
+  // Build a set of completed step names from `step:*` checkpoints with value '1'.
+  // The step runner writes these — `migration_*` registry keys are no longer
+  // written by migration functions (Phase 2 removed withCrashRecovery).
+  const completedStepNames = new Set(
     rows
-      .filter((r) => r.value !== "started" && r.value !== "rolling_back")
-      .map((r) => r.key),
+      .filter(
+        (r) =>
+          r.key.startsWith(STEP_CHECKPOINT_PREFIX) && r.value === "1",
+      )
+      .map((r) => r.key.slice(STEP_CHECKPOINT_PREFIX.length)),
   );
 
+  // Map registry entries to their step names to determine which migrations
+  // are completed. Multiple registry entries can share the same step name
+  // (e.g., 162's two entries both map to migrateGuardianTimestampsEpochMs).
+  const completed = new Set<string>();
+  for (const entry of MIGRATION_REGISTRY) {
+    if (completedStepNames.has(entry.stepName)) {
+      completed.add(entry.key);
+    }
+  }
+
   const dependencyViolations: Array<{
     migration: string;
     missingDependency: string;
@@ -216,21 +120,11 @@ export function validateMigrationState(
     );
   }
 
-  // Detect checkpoints that exist in the database but have no corresponding
-  // registry entry — these are from a newer version of the daemon.
-  //
-  // The memory_checkpoints table is a general-purpose key-value store also
-  // used by non-migration subsystems (e.g., "empty_state:greeting:text",
-  // "conversation_starters:item_count_at_last_gen"). Filter to only keys
-  // that follow migration naming conventions before comparing against the
-  // registry to avoid false-positive warnings.
-  const registryKeys = new Set(MIGRATION_REGISTRY.map((e) => e.key));
-  const isMigrationKey = (k: string): boolean =>
-    k.startsWith("migration_") ||
-    k.startsWith("backfill_") ||
-    k.startsWith("drop_");
-  const unknownCheckpoints = [...completed].filter(
-    (k) => isMigrationKey(k) && !registryKeys.has(k),
+  // Detect step checkpoints that exist in the database but have no
+  // corresponding registry entry — these are from a newer version of the daemon.
+  const registryStepNames = new Set(MIGRATION_REGISTRY.map((e) => e.stepName));
+  const unknownCheckpoints = [...completedStepNames].filter(
+    (name) => !registryStepNames.has(name),
   );
 
   if (unknownCheckpoints.length > 0) {
@@ -260,7 +154,9 @@ export function validateMigrationState(
  * **Checkpoint state**: Each rolled-back migration's checkpoint is deleted
  * from `memory_checkpoints`. If the process crashes mid-rollback, the
  * `"rolling_back"` marker is detected and cleared by
- * `recoverCrashedMigrations` on the next startup.
+ * `recoverCrashedMigrations` on the next startup. The forward-step checkpoints
+ * recorded by the migration runner (the `step:` namespace) for the rolled-back
+ * entries are also discarded so a later upgrade re-applies them.
  *
  * **Warning — data loss**: Some down() migrations may not fully restore the
  * original state (e.g., DROP TABLE migrations recreate the table but cannot
@@ -282,7 +178,7 @@ export function rollbackMemoryMigration(
 ): string[] {
   const raw = getSqliteFrom(database);
 
-  // Read completed checkpoints to determine which migrations have been applied.
+  // Read step checkpoints to determine which migrations have been applied.
   let rows: Array<{ key: string; value: string }>;
   try {
     rows = raw
@@ -292,36 +188,45 @@ export function rollbackMemoryMigration(
     return [];
   }
 
-  const completedKeys = new Set(
+  // Build a set of completed step names from `step:*` checkpoints with value '1'.
+  const completedStepNames = new Set(
     rows
-      .filter((r) => r.value !== "started" && r.value !== "rolling_back")
-      .map((r) => r.key),
+      .filter(
+        (r) =>
+          r.key.startsWith(STEP_CHECKPOINT_PREFIX) && r.value === "1",
+      )
+      .map((r) => r.key.slice(STEP_CHECKPOINT_PREFIX.length)),
   );
 
-  // Find registry entries with version > targetVersion that have completed checkpoints.
+  // Find registry entries with version > targetVersion whose step is completed.
+  // Deduplicate by stepName — multiple registry entries can share the same step
+  // (e.g., 162's two entries), and we only need to clear the step checkpoint once.
   const toRollback = MIGRATION_REGISTRY.filter(
-    (entry) => entry.version > targetVersion && completedKeys.has(entry.key),
+    (entry) =>
+      entry.version > targetVersion &&
+      completedStepNames.has(entry.stepName),
   ).sort((a, b) => b.version - a.version); // reverse version order
 
   const rolledBack: string[] = [];
 
   for (const entry of toRollback) {
+    const stepKey = `${STEP_CHECKPOINT_PREFIX}${entry.stepName}`;
+
     // Mark as rolling_back for crash recovery — if the process crashes here,
     // recoverCrashedMigrations will clear this checkpoint on next startup.
     raw
       .query(
-        `UPDATE memory_checkpoints SET value = 'rolling_back', updated_at = ? WHERE key = ?`,
+        `INSERT OR REPLACE INTO memory_checkpoints (key, value, updated_at) VALUES (?, 'rolling_back', ?)`,
       )
-      .run(Date.now(), entry.key);
+      .run(stepKey, Date.now());
 
     // Execute the down migration — let it manage its own transaction lifecycle.
     // Many down() functions call BEGIN/COMMIT internally or use PRAGMA statements
     // that are no-ops inside a transaction.
     entry.down(database);
 
-    // Delete the checkpoint after down() succeeds — outside any transaction
-    // so it's not affected by down()'s internal transaction management.
-    raw.query(`DELETE FROM memory_checkpoints WHERE key = ?`).run(entry.key);
+    // Delete the step checkpoint after down() succeeds.
+    raw.query(`DELETE FROM memory_checkpoints WHERE key = ?`).run(stepKey);
 
     log.info(
       { key: entry.key, version: entry.version },

package/src/memory/raw-query.ts

@@ -30,9 +30,9 @@
  * filtering, ordering, pagination — use Drizzle.
  */
 
-import type { SQLQueryBindings } from "bun:sqlite";
+import type { Database, SQLQueryBindings } from "bun:sqlite";
 
-import { getSqlite } from "./db-connection.js";
+import { getLogsSqlite, getMemorySqlite, getSqlite } from "./db-connection.js";
 
 type SqlParam = SQLQueryBindings;
 
@@ -82,6 +82,51 @@ export function rawChanges(): number {
   return (getSqlite().query("SELECT changes() AS c").get() as { c: number }).c;
 }
 
+// ---------------------------------------------------------------------------
+// Typed query helpers for the dedicated memory connection (assistant-memory.db)
+// ---------------------------------------------------------------------------
+
+/** The memory connection, or a thrown error when it cannot be opened. */
+function memorySqlite(): Database {
+  const sqlite = getMemorySqlite();
+  if (!sqlite) throw new Error("memory database unavailable");
+  return sqlite;
+}
+
+/** The logs connection, or a thrown error when it cannot be opened. */
+function logsSqlite(): Database {
+  const sqlite = getLogsSqlite();
+  if (!sqlite) throw new Error("logs database unavailable");
+  return sqlite;
+}
+
+/** {@link rawAll} against the memory connection. */
+export function rawMemoryAll<T>(sql: string, ...params: SqlParam[]): T[] {
+  return memorySqlite()
+    .query(sql)
+    .all(...params) as T[];
+}
+
+/** {@link rawRun} against the memory connection. */
+export function rawMemoryRun(sql: string, ...params: SqlParam[]): number {
+  const sqlite = memorySqlite();
+  sqlite.query(sql).run(...params);
+  return (sqlite.query("SELECT changes() AS c").get() as { c: number }).c;
+}
+
+/** {@link rawChanges} against the memory connection. */
+export function rawMemoryChanges(): number {
+  return (memorySqlite().query("SELECT changes() AS c").get() as { c: number })
+    .c;
+}
+
+/** {@link rawRun} against the logs connection. */
+export function rawLogsRun(sql: string, ...params: SqlParam[]): number {
+  const sqlite = logsSqlite();
+  sqlite.query(sql).run(...params);
+  return (sqlite.query("SELECT changes() AS c").get() as { c: number }).c;
+}
+
 /**
  * Delete all rows from the given tables in a single transaction.
  *

package/src/memory/skill-loaded-events-store.test.ts

@@ -22,7 +22,7 @@ import {
   recordSkillLoadedEvent,
 } from "./skill-loaded-events-store.js";
 
-initializeDb();
+await initializeDb();
 
 function insertEvent(
   id: string,