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

package/src/memory/migrations/schema-introspection.ts

@@ -16,3 +16,17 @@ export function tableHasColumn(
 
   return columns.some((column) => column.name === columnName);
 }
+
+/**
+ * True when a base table with the given name exists. Useful for migrations
+ * that recreate indexes — `CREATE INDEX IF NOT EXISTS` guards against the
+ * index already existing, but still throws "no such table" if the underlying
+ * table was already dropped by an earlier migration.
+ */
+export function tableExists(database: DrizzleDb, tableName: string): boolean {
+  const raw = getSqliteFrom(database);
+  const row = raw
+    .query(`SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?`)
+    .get(tableName);
+  return row != null;
+}

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/prompt-override.ts

@@ -0,0 +1,129 @@
+/**
+ * Shared loader for file-based prompt overrides.
+ *
+ * The memory router, consolidation, and v3 selector prompts each ship a bundled
+ * default that operators may replace by pointing a config field at a file. The
+ * load-with-fallback semantics are identical across all three, so they live
+ * here: path resolution, the size guard, and the permissive fallback that keeps
+ * retrieval/consolidation working when an override is missing or malformed.
+ *
+ * The caller owns the bundled default and any placeholder substitution — this
+ * module only decides whether a usable override file exists and returns its raw
+ * contents, or `null` to mean "fall back to the bundled prompt."
+ */
+
+import { lstatSync, readFileSync } from "node:fs";
+import { homedir } from "node:os";
+import { isAbsolute, join } from "node:path";
+
+/**
+ * Minimal logger surface the loader needs. Structurally compatible with
+ * `pino.Logger`, so callers pass their module logger directly and the warnings
+ * keep that logger's namespace; tests pass a lightweight recorder.
+ */
+export interface PromptOverrideLogger {
+  warn(obj: object, msg: string): void;
+}
+
+/**
+ * Hard upper bound on an override file. The bundled prompts are kilobytes; 1 MiB
+ * is generous headroom for a hand-edit while preventing a `settings.write`
+ * principal from pointing the field at a giant file (or an unsizeable stream
+ * `lstat` can't cap on its own) and slurping it into memory on every call.
+ */
+export const MAX_PROMPT_OVERRIDE_BYTES = 1 * 1024 * 1024;
+
+/**
+ * Resolve a configured override path to an absolute path: a leading `~/` expands
+ * to the home directory, an absolute path is used as-is, and a relative path
+ * resolves under `workspaceDir`.
+ */
+export function resolveOverridePath(
+  overridePath: string,
+  workspaceDir: string,
+): string {
+  if (overridePath.startsWith("~/")) {
+    return join(homedir(), overridePath.slice(2));
+  }
+  if (isAbsolute(overridePath)) return overridePath;
+  return join(workspaceDir, overridePath);
+}
+
+/**
+ * Load a prompt-override file, returning its raw contents when the override is
+ * present and usable, or `null` (after logging a warning describing why) when
+ * the caller should fall back to its bundled prompt. A nullish `overridePath`
+ * (`null` or `undefined`) returns `null` without touching the filesystem.
+ *
+ * Fallback is intentionally total — an absent override, a missing, non-regular,
+ * oversized, unreadable, or empty/whitespace-only file all degrade to `null`.
+ * Memory retrieval and consolidation must never break because of a bad
+ * override, so `undefined` (an unset config field) is treated as "no override"
+ * rather than crashing on path resolution.
+ *
+ * `label` names the prompt in the warning messages (e.g. `"router prompt"`).
+ */
+export function loadPromptOverride(opts: {
+  overridePath: string | null | undefined;
+  workspaceDir: string;
+  log: PromptOverrideLogger;
+  label: string;
+}): string | null {
+  const { overridePath, workspaceDir, log, label } = opts;
+  if (overridePath == null) return null;
+
+  const resolvedPath = resolveOverridePath(overridePath, workspaceDir);
+  let contents: string;
+  try {
+    const stat = lstatSync(resolvedPath);
+    if (!stat.isFile()) {
+      log.warn(
+        {
+          configuredPath: overridePath,
+          resolvedPath,
+          reason: "not_regular_file",
+          fallback: "bundled",
+        },
+        `${label} override is not a regular file; using bundled prompt`,
+      );
+      return null;
+    }
+    if (stat.size > MAX_PROMPT_OVERRIDE_BYTES) {
+      log.warn(
+        {
+          configuredPath: overridePath,
+          resolvedPath,
+          size: stat.size,
+          limit: MAX_PROMPT_OVERRIDE_BYTES,
+          reason: "oversized_override",
+          fallback: "bundled",
+        },
+        `${label} override exceeds size limit; using bundled prompt`,
+      );
+      return null;
+    }
+    contents = readFileSync(resolvedPath, "utf-8");
+  } catch (err) {
+    const code = (err as NodeJS.ErrnoException).code;
+    log.warn(
+      { configuredPath: overridePath, resolvedPath, code, fallback: "bundled" },
+      `${label} override unreadable; using bundled prompt`,
+    );
+    return null;
+  }
+
+  if (contents.trim().length === 0) {
+    log.warn(
+      {
+        configuredPath: overridePath,
+        resolvedPath,
+        reason: "empty_override",
+        fallback: "bundled",
+      },
+      `${label} override is empty; using bundled prompt`,
+    );
+    return null;
+  }
+
+  return contents;
+}

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,

package/src/memory/task-memory-cleanup.ts

@@ -1,8 +1,49 @@
 import { getLogger } from "../util/logger.js";
-import { rawGet, rawRun } from "./raw-query.js";
+import { rawAll, rawGet, rawMemoryRun, rawRun } from "./raw-query.js";
 
 const log = getLogger("task-memory-cleanup");
 
+/**
+ * Max id placeholders per `IN (...)` chunk when cancelling jobs by an id set
+ * resolved on the main connection. SQLite's default bound-variable limit is
+ * well above this, but chunking keeps a pathological conversation (tens of
+ * thousands of messages/segments) from building one giant statement.
+ */
+const ID_CHUNK_SIZE = 500;
+
+/**
+ * Fail every pending/running `memory_jobs` row whose `json_extract(payload,
+ * jsonPath)` matches one of `ids`. `memory_jobs` lives in the dedicated memory
+ * connection while the id sets are resolved on the main connection, so the
+ * cancellation can't be a single cross-DB subquery — bind the ids directly
+ * (chunked) instead. Returns the total rows affected.
+ */
+function cancelJobsByPayloadIds(
+  jsonPath: string,
+  ids: string[],
+  reason: string,
+  now: number,
+): number {
+  if (ids.length === 0) return 0;
+  let affected = 0;
+  for (let i = 0; i < ids.length; i += ID_CHUNK_SIZE) {
+    const chunk = ids.slice(i, i + ID_CHUNK_SIZE);
+    const placeholders = chunk.map(() => "?").join(", ");
+    affected += rawMemoryRun(
+      `UPDATE memory_jobs
+          SET status = 'failed',
+              last_error = ?,
+              updated_at = ?
+        WHERE status IN ('pending', 'running')
+          AND json_extract(payload, '${jsonPath}') IN (${placeholders})`,
+      reason,
+      now,
+      ...chunk,
+    );
+  }
+  return affected;
+}
+
 /**
  * Check whether a conversation belongs to a failed task run or failed
  * schedule run. Derived from durable storage (task_runs / cron_runs)
@@ -112,23 +153,19 @@ export function cancelPendingJobsForConversation(
   const now = Date.now();
   let total = 0;
 
+  // The id sets live on the main connection; `memory_jobs` lives on the memory
+  // connection. Resolve each set on main first, then cancel by bound ids on the
+  // memory connection (the conversationId-keyed update needs no main lookup).
+
   // Jobs keyed by messageId: embed_attachment
-  total += rawRun(
-    `UPDATE memory_jobs
-        SET status = 'failed',
-            last_error = ?,
-            updated_at = ?
-      WHERE status IN ('pending', 'running')
-        AND json_extract(payload, '$.messageId') IN (
-          SELECT id FROM messages WHERE conversation_id = ?
-        )`,
-    reason,
-    now,
+  const messageIds = rawAll<{ id: string }>(
+    `SELECT id FROM messages WHERE conversation_id = ?`,
     conversationId,
-  );
+  ).map((r) => r.id);
+  total += cancelJobsByPayloadIds("$.messageId", messageIds, reason, now);
 
   // Jobs keyed by conversationId: graph_extract, build_conversation_summary
-  total += rawRun(
+  total += rawMemoryRun(
     `UPDATE memory_jobs
         SET status = 'failed',
             last_error = ?,
@@ -141,36 +178,20 @@ export function cancelPendingJobsForConversation(
   );
 
   // Jobs keyed by segmentId: embed_segment (segments belong to the conversation)
-  total += rawRun(
-    `UPDATE memory_jobs
-        SET status = 'failed',
-            last_error = ?,
-            updated_at = ?
-      WHERE status IN ('pending', 'running')
-        AND json_extract(payload, '$.segmentId') IN (
-          SELECT id FROM memory_segments WHERE conversation_id = ?
-        )`,
-    reason,
-    now,
+  const segmentIds = rawAll<{ id: string }>(
+    `SELECT id FROM memory_segments WHERE conversation_id = ?`,
     conversationId,
-  );
+  ).map((r) => r.id);
+  total += cancelJobsByPayloadIds("$.segmentId", segmentIds, reason, now);
 
   // Jobs keyed by nodeId: embed_graph_node (nodes sourced from this conversation)
-  total += rawRun(
-    `UPDATE memory_jobs
-        SET status = 'failed',
-            last_error = ?,
-            updated_at = ?
-      WHERE status IN ('pending', 'running')
-        AND json_extract(payload, '$.nodeId') IN (
-          SELECT mgn.id
-            FROM memory_graph_nodes mgn, json_each(mgn.source_conversations) jc
-           WHERE jc.value = ?
-        )`,
-    reason,
-    now,
+  const nodeIds = rawAll<{ id: string }>(
+    `SELECT mgn.id
+       FROM memory_graph_nodes mgn, json_each(mgn.source_conversations) jc
+      WHERE jc.value = ?`,
     conversationId,
-  );
+  ).map((r) => r.id);
+  total += cancelJobsByPayloadIds("$.nodeId", nodeIds, reason, now);
 
   if (total > 0) {
     log.info(
@@ -192,7 +213,7 @@ function cancelPendingExtractionJobsForConversation(
   conversationId: string,
 ): number {
   const now = Date.now();
-  const cancelled = rawRun(
+  const cancelled = rawMemoryRun(
     `UPDATE memory_jobs
         SET status = 'failed',
             last_error = 'conversation_failed',

package/src/memory/tool-executed-events-store.test.ts

@@ -28,7 +28,7 @@ import { initializeDb } from "./db-init.js";
 import { conversations, toolInvocations } from "./schema.js";
 import { queryUnreportedToolExecutedEvents } from "./tool-executed-events-store.js";
 
-initializeDb();
+await initializeDb();
 
 const CONVERSATION_ID = "conv-tool-executed-store-test";
 

package/src/memory/turn-trace-store.test.ts

@@ -39,7 +39,7 @@ import {
   type TurnTraceBoundary,
 } from "./turn-trace-store.js";
 
-initializeDb();
+await initializeDb();
 
 function purge(): void {
   const db = getDb();