fix(db): preserve exact D2 rows during sync reconciliation

.changeset/sixty-nights-refuse.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/db': patch
+---
+
+Fix live queries throwing when updating an optimistically inserted row after synchronous sync confirmation.

packages/db/src/query/effect.ts

@@ -17,7 +17,7 @@ import {
   computeSubscriptionOrderByHints,
   extractCollectionSources,
   extractCollectionsFromQuery,
-  filterDuplicateInserts,
+  prepareChangesForD2,
   sendChangesToInput,
   splitUpdates,
   trackBiggestSentValue,
@@ -388,10 +388,10 @@ class EffectPipelineRunner<TRow extends object, TKey extends string | number> {
 
   // Subscription management
   private readonly unsubscribeCallbacks = new Set<() => void>()
-  // Duplicate insert prevention per lexical source
-  private readonly sentToD2KeysBySource = new Map<
+  // Exact rows contributed to D2 per lexical source
+  private readonly sentToD2RowsBySource = new Map<
     string,
-    Set<string | number>
+    Map<string | number, Record<string, unknown>>
   >()
 
   // Output accumulator
@@ -513,8 +513,8 @@ class EffectPipelineRunner<TRow extends object, TKey extends string | number> {
       const { sourceId, alias, collection } = source
       const collectionId = collection.id
 
-      // Initialise per-source duplicate tracking
-      this.sentToD2KeysBySource.set(sourceId, new Set())
+      // Initialise per-source D2 contribution tracking
+      this.sentToD2RowsBySource.set(sourceId, new Map())
 
       // Discover dependencies: if source collection is itself a live query
       // collection, its builder must run first during transaction flushes.
@@ -801,11 +801,10 @@ class EffectPipelineRunner<TRow extends object, TKey extends string | number> {
     const input = this.inputs[sourceId]
     if (!input) return 0
 
-    // Filter duplicates per lexical source
-    const sentKeys = this.sentToD2KeysBySource.get(sourceId)!
-    const filtered = filterDuplicateInserts(changes, sentKeys)
+    const sentRows = this.sentToD2RowsBySource.get(sourceId)!
+    const reconciled = prepareChangesForD2(changes, sentRows)
 
-    return sendChangesToInput(input, filtered)
+    return sendChangesToInput(input, reconciled)
   }
 
   /**
@@ -1040,11 +1039,11 @@ class EffectPipelineRunner<TRow extends object, TKey extends string | number> {
     changes: Array<ChangeMessage<any, string | number>>,
     comparator: (a: any, b: any) => number,
   ): void {
-    const sentKeys = this.sentToD2KeysBySource.get(sourceId) ?? new Set()
+    const sentRows = this.sentToD2RowsBySource.get(sourceId) ?? new Map()
     const result = trackBiggestSentValue(
       changes,
       this.biggestSentValue.get(sourceId),
-      sentKeys,
+      sentRows,
       comparator,
     )
     this.biggestSentValue.set(sourceId, result.biggest)
@@ -1069,7 +1068,7 @@ class EffectPipelineRunner<TRow extends object, TKey extends string | number> {
       }
     }
     this.unsubscribeCallbacks.clear()
-    this.sentToD2KeysBySource.clear()
+    this.sentToD2RowsBySource.clear()
     this.pendingChanges.clear()
     this.lazySources.clear()
     this.demand.clear()

packages/db/src/query/live/collection-subscriber.ts

@@ -5,7 +5,7 @@ import {
 import {
   computeOrderedLoadCursor,
   computeSubscriptionOrderByHints,
-  filterDuplicateInserts,
+  prepareChangesForD2,
   sendChangesToInput,
   splitUpdates,
   trackBiggestSentValue,
@@ -46,10 +46,9 @@ export class CollectionSubscriber<
     { resolve: () => void }
   >()
 
-  // Track keys that have been sent to the D2 pipeline to prevent duplicate inserts
-  // This is necessary because different code paths (initial load, change events)
-  // can potentially send the same item to D2 multiple times.
-  private sentToD2Keys = new Set<string | number>()
+  // Keep the exact value contributed for every source key. This both prevents
+  // duplicate inserts and ensures later updates retract the value currently in D2.
+  private sentToD2Rows = new Map<string | number, Record<string, unknown>>()
 
   // Direct load tracking callback for ordered path (set during subscribeToOrderedChanges,
   // used by loadNextItems for subsequent requestLimitedSnapshot calls)
@@ -237,16 +236,16 @@ export class CollectionSubscriber<
     callback?: () => boolean,
   ) {
     const changesArray = Array.isArray(changes) ? changes : [...changes]
-    const filteredChanges = filterDuplicateInserts(
+    const reconciledChanges = prepareChangesForD2(
       changesArray,
-      this.sentToD2Keys,
+      this.sentToD2Rows,
     )
 
     // currentSyncState and input are always defined when this method is called
     // (only called from active subscriptions during a sync session)
     const input =
       this.collectionConfigBuilder.currentSyncState!.inputs[this.sourceId]!
-    const sentChanges = sendChangesToInput(input, filteredChanges)
+    const sentChanges = sendChangesToInput(input, reconciledChanges)
 
     // Do not provide the callback that loads more data
     // if there's no more data to load
@@ -351,14 +350,14 @@ export class CollectionSubscriber<
     subscriptionHolder.current = subscription
     this.registerSubscriptionCleanup(subscription)
 
-    // Listen for truncate events to reset cursor tracking state and sentToD2Keys
+    // Listen for truncate events to reset cursor and D2 source tracking state.
     // This ensures that after a must-refetch/truncate, we don't use stale cursor data
     // and allow re-inserts of previously sent keys
     const truncateUnsubscribe = this.collection.on(`truncate`, () => {
       this.biggest = undefined
       this.lastLoadRequestKey = undefined
       this.pendingOrderedLoadPromise = undefined
-      this.sentToD2Keys.clear()
+      this.sentToD2Rows.clear()
     })
 
     // Clean up truncate listener when subscription is unsubscribed
@@ -542,7 +541,7 @@ export class CollectionSubscriber<
     const result = trackBiggestSentValue(
       changes,
       this.biggest,
-      this.sentToD2Keys,
+      this.sentToD2Rows,
       comparator,
     )
     this.biggest = result.biggest

packages/db/src/query/live/utils.ts

@@ -140,29 +140,40 @@ export function* splitUpdates<
 }
 
 /**
- * Filter changes to prevent duplicate inserts to a D2 pipeline.
- * Maintains D2 multiplicity at 1 for visible items so that deletes
- * properly reduce multiplicity to 0.
+ * Prepare changes for a D2 pipeline by preventing duplicate inserts and
+ * reconciling update/delete retractions with previously sent rows.
+ * Maintains D2 multiplicity at 1 for visible items and ensures retractions
+ * exactly match the row previously contributed for each key.
  *
- * Mutates `sentKeys` in place: adds keys on insert, removes on delete.
+ * Mutates `sentRows` in place: records rows on insert/update, removes them
+ * on delete.
  */
-export function filterDuplicateInserts(
-  changes: Array<ChangeMessage<any, string | number>>,
-  sentKeys: Set<string | number>,
-): Array<ChangeMessage<any, string | number>> {
-  const filtered: Array<ChangeMessage<any, string | number>> = []
-  for (const change of changes) {
+export function prepareChangesForD2(
+  changes: Array<ChangeMessage<Record<string, unknown>, string | number>>,
+  sentRows: Map<string | number, Record<string, unknown>>,
+): Array<ChangeMessage<Record<string, unknown>, string | number>> {
+  return changes.flatMap((change) => {
     if (change.type === `insert`) {
-      if (sentKeys.has(change.key)) {
-        continue // Skip duplicate
-      }
-      sentKeys.add(change.key)
-    } else if (change.type === `delete`) {
-      sentKeys.delete(change.key)
+      if (sentRows.has(change.key)) return []
+      sentRows.set(change.key, change.value)
+      return [change]
     }
-    filtered.push(change)
-  }
-  return filtered
+
+    const previousValue = sentRows.get(change.key)
+    if (change.type === `delete`) {
+      sentRows.delete(change.key)
+      return [
+        previousValue === undefined
+          ? change
+          : { ...change, value: previousValue },
+      ]
+    }
+
+    sentRows.set(change.key, change.value)
+    return [
+      previousValue === undefined ? change : { ...change, previousValue },
+    ]
+  })
 }
 
 /**
@@ -172,15 +183,19 @@ export function filterDuplicateInserts(
  *
  * @param changes   - changes to process (deletes are skipped)
  * @param current   - the current biggest value (or undefined if none)
- * @param sentKeys  - set of keys already sent to D2 (for new-key detection)
+ * @param sentKeys  - lookup of keys already sent to D2 (for new-key detection)
  * @param comparator - orderBy comparator
  * @returns `{ biggest, shouldResetLoadKey }` — the new biggest value and
  *          whether the caller should clear its last-load-request-key
  */
+interface SentKeyLookup {
+  has: (key: string | number) => boolean
+}
+
 export function trackBiggestSentValue(
   changes: Array<ChangeMessage<any, string | number>>,
   current: unknown | undefined,
-  sentKeys: Set<string | number>,
+  sentKeys: SentKeyLookup,
   comparator: (a: any, b: any) => number,
 ): { biggest: unknown; shouldResetLoadKey: boolean } {
   let biggest = current

packages/db/tests/collection-subscriber-duplicate-inserts.test.ts

@@ -2,8 +2,9 @@ import { describe, expect, it } from 'vitest'
 import { createCollection } from '../src/collection/index.js'
 import { BTreeIndex } from '../src/indexes/btree-index.js'
 import { createLiveQueryCollection, eq } from '../src/query/index.js'
+import { prepareChangesForD2 } from '../src/query/live/utils.js'
 import { mockSyncCollectionOptions } from './utils.js'
-import type { ChangeMessage } from '../src/types.js'
+import type { ChangeMessage, SyncConfig } from '../src/types.js'
 
 /**
  * Tests for duplicate insert prevention in the D2 pipeline.
@@ -40,6 +41,87 @@ type Order = {
 }
 
 describe(`CollectionSubscriber duplicate insert prevention`, () => {
+  it(`retracts the exact source row previously contributed for a key`, () => {
+    const sentRows = new Map<string | number, Record<string, unknown>>()
+    const inserted = { id: `1`, status: `draft` }
+    const changed = { id: `1`, status: `published` }
+
+    prepareChangesForD2(
+      [{ type: `insert`, key: `1`, value: inserted }],
+      sentRows,
+    )
+    const reconciled = prepareChangesForD2(
+      [
+        {
+          type: `update`,
+          key: `1`,
+          value: changed,
+          previousValue: changed,
+        },
+      ],
+      sentRows,
+    )
+
+    expect(reconciled).toEqual([
+      {
+        type: `update`,
+        key: `1`,
+        value: changed,
+        previousValue: inserted,
+      },
+    ])
+  })
+
+  it(`does not throw when updating an optimistically inserted row confirmed by sync and observed by a live query`, async () => {
+    type Row = { id: string; title: string }
+    let sync: Parameters<SyncConfig<Row>[`sync`]>[0] | undefined
+
+    const collection = createCollection<Row>({
+      getKey: (row) => row.id,
+      sync: {
+        sync: (config) => {
+          sync = config
+          config.markReady()
+        },
+      },
+      // eslint-disable-next-line @typescript-eslint/require-await -- onInsert must return a Promise
+      onInsert: async ({ transaction }) => {
+        if (!sync) throw new Error(`Sync config was not initialized`)
+        sync.begin()
+        sync.write({
+          type: `insert`,
+          value: transaction.mutations[0].modified,
+        })
+        sync.commit()
+      },
+      onUpdate: async () => {},
+    })
+
+    const query = createLiveQueryCollection((q) =>
+      q.from({ row: collection }).select(({ row }) => ({
+        id: row.id,
+        title: row.title,
+      })),
+    )
+
+    const subscription = query.subscribeChanges(() => {})
+
+    await collection.insert({ id: `1`, title: `one` }).isPersisted.promise
+
+    expect(() => {
+      collection.update(`1`, (draft) => {
+        draft.title = `two`
+      })
+    }).not.toThrow()
+
+    expect(collection.size).toBe(1)
+    expect(collection.get(`1`)).toMatchObject({ id: `1`, title: `two` })
+    expect(query.size).toBe(1)
+    expect(query.get(`1`)).toMatchObject({ id: `1`, title: `two` })
+
+    subscription.unsubscribe()
+  })
+
   it(`should properly delete items from live query with orderBy + limit`, async () => {
     // This test verifies that items can be properly deleted from a live query
     // with orderBy + limit. If duplicate inserts reach D2, the delete won't work.