mason: fix plugin integrity races

Co-authored-by: Alfonso [Magic Context] <288211368+alfonso-magic-context@users.noreply.github.com>

Author: ualtinok

packages/plugin/src/features/magic-context/memory/storage-memory-embeddings.ts

@@ -16,6 +16,7 @@ interface StoredModelIdRow {
 }
 
 const saveEmbeddingStatements = new WeakMap<Database, PreparedStatement>();
+const saveEmbeddingIfHashMatchesStatements = new WeakMap<Database, PreparedStatement>();
 const loadAllEmbeddingsStatements = new WeakMap<Database, PreparedStatement>();
 const deleteEmbeddingStatements = new WeakMap<Database, PreparedStatement>();
 const getStoredModelIdStatements = new WeakMap<Database, PreparedStatement>();
@@ -57,6 +58,17 @@ function getSaveEmbeddingStatement(db: Database): PreparedStatement {
     return stmt;
 }
 
+function getSaveEmbeddingIfHashMatchesStatement(db: Database): PreparedStatement {
+    let stmt = saveEmbeddingIfHashMatchesStatements.get(db);
+    if (!stmt) {
+        stmt = db.prepare(
+            "INSERT INTO memory_embeddings (memory_id, embedding, model_id) SELECT ?, ?, ? FROM memories WHERE id = ? AND normalized_hash = ? ON CONFLICT(memory_id, model_id) DO UPDATE SET embedding = excluded.embedding",
+        );
+        saveEmbeddingIfHashMatchesStatements.set(db, stmt);
+    }
+    return stmt;
+}
+
 function getLoadAllEmbeddingsStatement(db: Database): PreparedStatement {
     let stmt = loadAllEmbeddingsStatements.get(db);
     if (!stmt) {
@@ -131,6 +143,28 @@ export function saveEmbedding(
     getSaveEmbeddingStatement(db).run(memoryId, blob, modelId);
 }
 
+/** Save an embedding only if the memory row still has the same normalized hash
+ *  we embedded. If the content changed while the provider call was in flight,
+ *  the stale vector is discarded instead of resurrecting an out-of-date row. */
+export function saveEmbeddingIfHashMatches(
+    db: Database,
+    memoryId: number,
+    embedding: Float32Array,
+    modelId: string,
+    normalizedHash: string,
+): boolean {
+    const blob = Buffer.from(embedding.buffer, embedding.byteOffset, embedding.byteLength);
+    return (
+        getSaveEmbeddingIfHashMatchesStatement(db).run(
+            memoryId,
+            blob,
+            modelId,
+            memoryId,
+            normalizedHash,
+        ).changes > 0
+    );
+}
+
 export function loadAllEmbeddings(
     db: Database,
     projectPath: string,

packages/plugin/src/features/magic-context/memory/storage-memory.ts

@@ -117,6 +117,11 @@ export interface MemoryCountsByStatus {
     mergedIds: number[];
 }
 
+export interface InsertMemoryResult {
+    memory: Memory;
+    inserted: boolean;
+}
+
 interface MemoryCountByStatusRow {
     id: number;
     status: MemoryStatus;
@@ -221,6 +226,14 @@ function isVerificationStatus(value: unknown): value is VerificationStatus {
     return typeof value === "string" && value in VERIFICATION_STATUS_LOOKUP;
 }
 
+function isUniqueConstraintError(error: unknown): boolean {
+    return (
+        error instanceof Error &&
+        "code" in error &&
+        (error as { code?: unknown }).code === "SQLITE_CONSTRAINT_UNIQUE"
+    );
+}
+
 function isNullableString(value: unknown): value is string | null {
     return value === null || typeof value === "string";
 }
@@ -500,16 +513,19 @@ function getMemoryCountsByStatusStatement(db: Database): PreparedStatement {
     return stmt;
 }
 
-export function insertMemory(db: Database, input: MemoryInput): Memory {
-    const now = Date.now();
-    const normalizedHash = computeNormalizedHash(input.content);
+function buildInsertMemoryValues(
+    input: MemoryInput,
+    normalizedHash: string,
+    now: number,
+    includeImportance: boolean,
+): Array<string | number | null> {
     const insertValues: Array<string | number | null> = [
         input.projectPath,
         input.category,
         input.content,
         normalizedHash,
     ];
-    if (hasMemoryImportanceColumn(db)) {
+    if (includeImportance) {
         insertValues.push(input.importance ?? 50);
     }
     insertValues.push(
@@ -530,18 +546,61 @@ export function insertMemory(db: Database, input: MemoryInput): Memory {
         null,
         input.metadataJson ?? null,
     );
-    const result = getInsertMemoryStatement(db).run(...insertValues);
+    return insertValues;
+}
 
-    const insertedResult = result as { lastInsertRowid?: number | bigint };
-    const inserted = getMemoryById(db, Number(insertedResult.lastInsertRowid));
+function loadInsertedMemory(db: Database, rowid: number | bigint | undefined): Memory {
+    const inserted = getMemoryById(db, Number(rowid));
     if (!inserted) {
         throw new Error("Failed to load inserted memory row");
     }
+    return inserted;
+}
+
+export function insertMemory(db: Database, input: MemoryInput): Memory {
+    const now = Date.now();
+    const normalizedHash = computeNormalizedHash(input.content);
+    const insertValues = buildInsertMemoryValues(
+        input,
+        normalizedHash,
+        now,
+        hasMemoryImportanceColumn(db),
+    );
+    const result = getInsertMemoryStatement(db).run(...insertValues);
+
+    const insertedResult = result as { lastInsertRowid?: number | bigint };
+    const inserted = loadInsertedMemory(db, insertedResult.lastInsertRowid);
 
     invalidateProject(input.projectPath);
     return inserted;
 }
 
+/**
+ * Shared-DB callers can race between their exact-hash pre-check and INSERT. When
+ * the unique constraint wins, treat it as the same exact-dedup path the tool uses:
+ * bump seen_count/last_seen_at on the existing row and return it instead of
+ * surfacing a transient write failure.
+ */
+export function insertMemoryIdempotent(db: Database, input: MemoryInput): InsertMemoryResult {
+    try {
+        return { memory: insertMemory(db, input), inserted: true };
+    } catch (error) {
+        if (!isUniqueConstraintError(error)) {
+            throw error;
+        }
+        const normalizedHash = computeNormalizedHash(input.content);
+        const existing = getMemoryByHash(db, input.projectPath, input.category, normalizedHash);
+        if (!existing) {
+            throw error;
+        }
+        updateMemorySeenCount(db, existing.id);
+        return {
+            memory: getMemoryById(db, existing.id) ?? existing,
+            inserted: false,
+        };
+    }
+}
+
 export function getMemoryByHash(
     db: Database,
     projectPath: string,

packages/plugin/src/features/magic-context/message-index-async.test.ts

@@ -49,6 +49,16 @@ function countMessageRows(db: Database, sessionId: string, messageId: string): n
     return typeof row?.count === "number" ? row.count : 0;
 }
 
+function searchMessageIds(db: Database, sessionId: string, ftsQuery: string): string[] {
+    return (
+        db
+            .prepare(
+                "SELECT message_id FROM message_history_fts WHERE session_id = ? AND message_history_fts MATCH ? ORDER BY bm25(message_history_fts), CAST(message_ordinal AS INTEGER) ASC",
+            )
+            .all(sessionId, ftsQuery) as Array<{ message_id: string }>
+    ).map((row) => row.message_id);
+}
+
 describe("message-index-async", () => {
     let db: Database;
 
@@ -92,6 +102,50 @@ describe("message-index-async", () => {
         expect(countMessageRows(db, "ses-overlap", "m-1")).toBe(1);
     });
 
+    it("reconciles a failed incremental hole even after a later incremental success advanced the watermark", async () => {
+        const originalPrepare = db.prepare.bind(db);
+        let failMessageId: string | null = "m-2";
+        (db as unknown as { prepare: typeof db.prepare }).prepare = ((sql: string) => {
+            const stmt = originalPrepare(sql);
+            if (sql.startsWith("INSERT INTO message_history_fts")) {
+                const run = stmt.run.bind(stmt);
+                (stmt as unknown as { run: typeof stmt.run }).run = ((...args: unknown[]) => {
+                    if (failMessageId !== null && args[2] === failMessageId) {
+                        throw new Error("synthetic incremental failure");
+                    }
+                    return run(...(args as Parameters<typeof stmt.run>));
+                }) as typeof stmt.run;
+            }
+            return stmt;
+        }) as typeof db.prepare;
+
+        const fullHistory = [
+            message("m-1", 1, "alpha indexed first"),
+            message("m-2", 2, "beta hole should come back"),
+            message("m-3", 3, "gamma later incremental succeeds"),
+        ];
+        scheduleReconciliation(db, "ses-hole", () => [fullHistory[0]!]);
+        await wait(20);
+        expect(isSessionReconciled("ses-hole")).toBe(true);
+
+        scheduleIncrementalIndex(db, "ses-hole", "m-2", () => fullHistory[1] ?? null);
+        await wait(140);
+        expect(countMessageRows(db, "ses-hole", "m-2")).toBe(0);
+        expect(isSessionReconciled("ses-hole")).toBe(false);
+
+        failMessageId = null;
+        scheduleIncrementalIndex(db, "ses-hole", "m-3", () => fullHistory[2] ?? null);
+        await wait(140);
+        expect(countMessageRows(db, "ses-hole", "m-3")).toBe(1);
+
+        scheduleReconciliation(db, "ses-hole", () => fullHistory);
+        await wait(20);
+
+        expect(searchMessageIds(db, "ses-hole", "beta")).toEqual(["m-2"]);
+        expect(countMessageRows(db, "ses-hole", "m-3")).toBe(1);
+        expect(isSessionReconciled("ses-hole")).toBe(true);
+    });
+
     it("clears and rebuilds after a removed message", async () => {
         const first = [message("m-1", 1, "old"), message("m-2", 2, "keep")];
         scheduleReconciliation(db, "ses-clear", () => first);

packages/plugin/src/features/magic-context/message-index-async.ts

@@ -6,6 +6,7 @@ import {
     getLastIndexedOrdinal,
     indexMessagesAfterOrdinal,
     indexSingleMessage,
+    markMessageIndexDirty,
 } from "./message-index";
 
 /**
@@ -47,8 +48,9 @@ function isDatabaseLockedError(error: unknown): boolean {
  *    the same `(sessionId,messageId)` inside 100ms are dropped.
  * 2. Per-session lazy reconciliation: the first transform/hook touch schedules
  *    one catch-up pass. It reads raw messages, resumes from
- *    `message_history_index.last_indexed_ordinal`, inserts missing newer rows,
- *    and advances the watermark to `messages.length`.
+ *    `message_history_index.last_indexed_ordinal`, rewinds from any recorded
+ *    dirty floor left by a failed incremental write, inserts missing rows, and
+ *    advances the watermark to `messages.length`.
  * 3. Revert/delete handling: `message.removed` clears all FTS rows + the
  *    watermark and then re-runs reconciliation. Searches during that rebuild
  *    window correctly see no message hits.
@@ -160,10 +162,13 @@ export function scheduleReconciliation(
     reconciliationScheduledSessions.add(sessionId);
 
     defer(() => {
-        void reconcileSessionIndex(db, sessionId, readMessages).catch((error) => {
-            reconciliationScheduledSessions.delete(sessionId);
-            logIndexingError(sessionId, "reconciliation", error);
-        });
+        void reconcileSessionIndex(db, sessionId, readMessages)
+            .catch((error) => {
+                logIndexingError(sessionId, "reconciliation", error);
+            })
+            .finally(() => {
+                reconciliationScheduledSessions.delete(sessionId);
+            });
     });
 }
 
@@ -181,14 +186,22 @@ export function scheduleIncrementalIndex(
     const timer = setTimeout(() => {
         incrementalTimers.delete(key);
         pendingIncrementalKeys.add(key);
+        let attemptedOrdinal: number | null = null;
         void runWithSessionLock(sessionId, () => {
             const message = readSingleMessage(sessionId, messageId);
             if (!message) {
                 return;
             }
+            attemptedOrdinal = message.ordinal;
             indexSingleMessage(db, sessionId, message);
         })
             .catch((error) => {
+                markMessageIndexDirty(
+                    db,
+                    sessionId,
+                    attemptedOrdinal ?? getLastIndexedOrdinal(db, sessionId) + 1,
+                );
+                reconciledSessions.delete(sessionId);
                 logIndexingError(sessionId, `incremental index for ${messageId}`, error);
             })
             .finally(() => {