diff --git a/packages/plugin/src/features/magic-context/project-embedding-registry.test.ts b/packages/plugin/src/features/magic-context/project-embedding-registry.test.ts index 6bb1e25db..969fe6e28 100644 --- a/packages/plugin/src/features/magic-context/project-embedding-registry.test.ts +++ b/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 { + 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 { + 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 { + 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); diff --git a/packages/plugin/src/features/magic-context/project-embedding-registry.ts b/packages/plugin/src/features/magic-context/project-embedding-registry.ts index 1abc241e0..01e17a42b 100644 --- a/packages/plugin/src/features/magic-context/project-embedding-registry.ts +++ b/packages/plugin/src/features/magic-context/project-embedding-registry.ts @@ -1981,6 +1981,14 @@ async function processShadowQueueItem(item: ShadowQueueItem): Promise { @@ -2035,6 +2043,11 @@ async function processShadowQueueItem(item: ShadowQueueItem): Promise { for (const row of rows) { @@ -2113,6 +2126,11 @@ async function processShadowQueueItem(item: ShadowQueueItem): Promise { 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; diff --git a/packages/plugin/src/features/magic-context/shadow-backfill-state.ts b/packages/plugin/src/features/magic-context/shadow-backfill-state.ts index 111e44926..e978d2b6a 100644 --- a/packages/plugin/src/features/magic-context/shadow-backfill-state.ts +++ b/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":