fix(plugin): discard shadow vectors when the registration is retired mid-embed

packages/plugin/src/features/magic-context/project-embedding-registry.test.ts

@@ -45,6 +45,7 @@ import {
     getEmbeddingCoverageStatus,
     getProjectEmbeddingSnapshot,
     getShadowBackfillStopReason,
+    getShadowBackfillRemaining,
     getShadowEmbeddingMeasurementCohort,
     markProjectLoadUntrusted,
     registerProjectEmbedding,
@@ -1991,6 +1992,168 @@ describe("project embedding registry", () => {
         expect(loadAllEmbeddings(db, projectIdentity, repeated!.modelId).size).toBe(3);
     });
 
+    it("discards shadow vectors when the registration is retired while the provider call is in flight", async () => {
+        const db = useTempDb();
+        const projectIdentity = "git:shadow-retired-in-flight";
+        registerProjectEmbedding(
+            db,
+            projectIdentity,
+            localConfig("model-primary"),
+            { memoryEnabled: true, gitCommitEnabled: false },
+            "/tmp/shadow-retired-in-flight",
+        );
+        for (let i = 0; i < 3; i++) {
+            const memory = insertMemory(db, {
+                projectPath: projectIdentity,
+                category: "CONSTRAINTS",
+                content: `retired shadow memory ${i}`,
+            });
+            saveEmbedding(db, memory.id, new Float32Array([i, 1]), currentModelId(projectIdentity));
+        }
+        const shadowConfig = {
+            provider: "synapse",
+            model: "synapse-model",
+            synapse_fingerprint: "fp-retired-in-flight",
+        } as unknown as EmbeddingConfig;
+        // Retire the shadow from inside the provider call. processShadowQueueItem
+        // captures the registration before awaiting, so this reproduces the
+        // unregister-during-embed window: without a post-await re-check the
+        // returned vectors are written under a registration that no longer exists.
+        _setTestProviderFactoryForProject(
+            () =>
+                new (class extends FakeEmbeddingProvider {
+                    override async embedBatch(texts: string[]): Promise<Float32Array[]> {
+                        unregisterProjectShadowEmbedding(projectIdentity);
+                        return texts.map((text) => new Float32Array([text.length, this.modelId.length]));
+                    }
+                })("shadow"),
+        );
+        const registration = registerProjectShadowEmbedding(
+            db,
+            projectIdentity,
+            shadowConfig,
+            "/tmp/shadow-retired-in-flight",
+        );
+        await flushShadowEmbeddingBacklog(projectIdentity);
+
+        expect(loadAllEmbeddings(db, projectIdentity, registration!.modelId).size).toBe(0);
+    });
+
+    it("discards commit shadow vectors when the registration is retired during the embed call", async () => {
+        const db = useTempDb();
+        const projectIdentity = "git:shadow-retired-commit-in-flight";
+        registerProjectEmbedding(
+            db,
+            projectIdentity,
+            localConfig("model-primary"),
+            { memoryEnabled: false, gitCommitEnabled: true },
+            "/tmp/shadow-retired-commit-in-flight",
+        );
+        const commits = ["c-a", "c-b", "c-c"].map((seed) => makeGitCommit(seed, 1000));
+        upsertCommits(db, projectIdentity, commits);
+        const primaryModelId = currentModelId(projectIdentity);
+        for (const commit of commits) {
+            saveCommitEmbedding(db, commit.sha, new Float32Array([1, 1]), primaryModelId);
+        }
+        _setTestProviderFactoryForProject(
+            () =>
+                new (class extends FakeEmbeddingProvider {
+                    override async embedBatch(texts: string[]): Promise<Float32Array[]> {
+                        unregisterProjectShadowEmbedding(projectIdentity);
+                        return texts.map(
+                            (text) => new Float32Array([text.length, this.modelId.length]),
+                        );
+                    }
+                })("shadow"),
+        );
+        const registration = registerProjectShadowEmbedding(
+            db,
+            projectIdentity,
+            {
+                provider: "synapse",
+                model: "synapse-model",
+                synapse_fingerprint: "fp-retired-commit",
+            } as unknown as EmbeddingConfig,
+            "/tmp/shadow-retired-commit-in-flight",
+        );
+        await flushShadowEmbeddingBacklog(projectIdentity);
+
+        expect(countEmbeddedCommits(db, projectIdentity, registration!.modelId)).toBe(0);
+    });
+
+    it("discards chunk shadow vectors when the registration is retired during the embed call", async () => {
+        const db = useTempDb();
+        const projectIdentity = "git:shadow-retired-chunk-in-flight";
+        const sessionId = "ses-shadow-retired-chunk-in-flight";
+        registerProjectEmbedding(
+            db,
+            projectIdentity,
+            localConfig("model-primary", 512),
+            { memoryEnabled: true, gitCommitEnabled: false },
+            "/tmp/shadow-retired-chunk-in-flight",
+        );
+        recordSessionProjectIdentity(db, sessionId, projectIdentity);
+        appendCompartments(db, sessionId, [
+            {
+                sequence: 0,
+                startMessage: 1,
+                endMessage: 1,
+                startMessageId: "a1",
+                endMessageId: "a1",
+                title: "Retired chunk",
+                content: "large transcript",
+                p1: "large transcript",
+            },
+        ]);
+        const content = Array.from({ length: 2_000 }, (_, index) => `token-${index}`).join(" ");
+        const ftsRow = db
+            .prepare(
+                "INSERT INTO message_history_fts (session_id, message_ordinal, message_id, role, content) VALUES (?, 1, 'a1', 'assistant', ?)",
+            )
+            .run(sessionId, content) as { lastInsertRowid: number | bigint };
+        recordMessageFtsRowid(db, sessionId, 1, ftsRow.lastInsertRowid);
+        _setTestProviderFactoryForProject((config) =>
+            config.provider === "synapse"
+                ? // Only the shadow lane retires itself mid-embed; the primary lane
+                  // must embed normally so the shadow has a row to mirror.
+                  new (class extends FakeEmbeddingProvider {
+                      override async embedBatch(texts: string[]): Promise<Float32Array[]> {
+                          unregisterProjectShadowEmbedding(projectIdentity);
+                          return texts.map(
+                              (text) => new Float32Array([text.length, this.modelId.length]),
+                          );
+                      }
+                  })(config.model)
+                : new FakeEmbeddingProvider(config.model),
+        );
+        // Seed the primary chunk lane so the shadow lane has a row to mirror.
+        expect(await embedUnembeddedCompartmentChunksForProject(db, projectIdentity, 8)).toBe(1);
+
+        const registration = registerProjectShadowEmbedding(
+            db,
+            projectIdentity,
+            {
+                provider: "synapse",
+                model: "synapse-model",
+                synapse_fingerprint: "fp-retired-chunk",
+            } as unknown as EmbeddingConfig,
+            "/tmp/shadow-retired-chunk-in-flight",
+        );
+        // Positive control: the fixture must produce a chunk candidate, otherwise
+        // "no rows written" would hold for the wrong reason.
+        expect(getShadowBackfillRemaining(db, projectIdentity).chunk).toBe(1);
+
+        await flushShadowEmbeddingBacklog(projectIdentity);
+
+        expect(
+            countRows(
+                db,
+                "SELECT COUNT(*) AS count FROM compartment_chunk_embeddings WHERE model_id = ?",
+                registration!.chunkModelId,
+            ),
+        ).toBe(0);
+    });
+
     it("does not dispose a shadow provider that is the same instance as the primary", async () => {
         const shared = new FakeEmbeddingProvider("shared");
         _setTestProviderFactoryForProject(() => shared);

packages/plugin/src/features/magic-context/project-embedding-registry.ts

@@ -1981,6 +1981,14 @@ async function processShadowQueueItem(item: ShadowQueueItem): Promise<ShadowBack
             db,
             "memory",
         );
+        // The provider call above can take seconds (or minutes on a cold model
+        // load). unregisterProjectShadowEmbedding may have retired this shadow
+        // while we were waiting, so re-check the live registration before
+        // writing any vectors.
+        const live = shadowRegistrations.get(item.projectIdentity);
+        if (!live || live.generation !== registration.generation) {
+            return { writes: 0, refusalReason: "registration_retired_during_embed" };
+        }
         let writes = 0;
         let hashGuardRejected = false;
         db.transaction(() => {
@@ -2035,6 +2043,11 @@ async function processShadowQueueItem(item: ShadowQueueItem): Promise<ShadowBack
             db,
             "commit",
         );
+        // Re-check after the provider round-trip; see the memory scope above.
+        const live = shadowRegistrations.get(item.projectIdentity);
+        if (!live || live.generation !== registration.generation) {
+            return { writes: 0, refusalReason: "registration_retired_during_embed" };
+        }
         let writes = 0;
         db.transaction(() => {
             for (const row of rows) {
@@ -2113,6 +2126,11 @@ async function processShadowQueueItem(item: ShadowQueueItem): Promise<ShadowBack
         })),
     );
     const embedded = await embedShadowItems(registration, items, db, "chunk");
+    // Re-check after the provider round-trip; see the memory scope above.
+    const live = shadowRegistrations.get(item.projectIdentity);
+    if (!live || live.generation !== registration.generation) {
+        return { writes: 0, refusalReason: "registration_retired_during_embed" };
+    }
     let writes = 0;
     let partialVectorSet = false;
     for (const item of prepared) {
@@ -2182,14 +2200,28 @@ async function runShadowWorker(): Promise<void> {
             shadowQueue.unshift(item);
             break;
         }
+        // A worker item may outlive its registration: retirement or a re-arm can
+        // land while the provider call is in flight. Publishing the outcome then
+        // re-creates state that retirement just cleared, and the stall detector
+        // reads that map — so a retired item's refusal could be attributed to a
+        // freshly re-armed registration and stop its backfill. Only record the
+        // outcome when the registration this item started under is still live.
+        const generationAtStart = shadowRegistrations.get(item.projectIdentity)?.generation;
+        const isStillCurrent = (): boolean =>
+            generationAtStart !== undefined &&
+            shadowRegistrations.get(item.projectIdentity)?.generation === generationAtStart;
         try {
             const outcome = await processShadowQueueItem(item);
-            shadowBackfillLastWriteOutcomes.set(`${item.projectIdentity}:${item.scope}`, outcome);
+            if (isStillCurrent()) {
+                shadowBackfillLastWriteOutcomes.set(`${item.projectIdentity}:${item.scope}`, outcome);
+            }
         } catch (error) {
-            shadowBackfillLastWriteOutcomes.set(`${item.projectIdentity}:${item.scope}`, {
-                writes: 0,
-                refusalReason: "provider_returned_no_vectors",
-            });
+            if (isStillCurrent()) {
+                shadowBackfillLastWriteOutcomes.set(`${item.projectIdentity}:${item.scope}`, {
+                    writes: 0,
+                    refusalReason: "provider_returned_no_vectors",
+                });
+            }
             log("[magic-context] Synapse shadow write failed:", error);
         }
         processed += item.ids.length;

packages/plugin/src/features/magic-context/shadow-backfill-state.ts

@@ -8,6 +8,7 @@ export type ShadowBackfillWriteRefusalReason =
     | "provider_returned_no_vectors"
     | "memory_hash_guard_rejected"
     | "candidate_rows_changed"
+    | "registration_retired_during_embed"
     | "chunk_fts_mapping_incomplete"
     | "chunk_empty_canonical_text"
     | "chunk_partial_vector_set"
@@ -75,6 +76,8 @@ export function describeShadowBackfillWriteRefusal(
             return "the memory normalized-hash guard rejected vectors because content changed in flight";
         case "candidate_rows_changed":
             return "the selected source rows changed before the writer loaded them";
+        case "registration_retired_during_embed":
+            return "the shadow registration was retired or replaced while the provider call was in flight, so the vectors were discarded";
         case "chunk_fts_mapping_incomplete":
             return "the chunk writer refused rows whose transcript ordinals are not fully mapped in FTS";
         case "chunk_empty_canonical_text":