fix: Fix broken array patching and ensure all values in draft Maps/Sets are finalized (#1201)

* Remove hardcoded key_ field name check

* Only build for `test:build` command

* Rewrite Vitest prod build config to ensure prod build is used

* Add tests for prod patch paths

* Ensure plain values in draft maps/sets get finalized

* Work around an error with prod tests

Somehow was seeing an error with `each` "not existing" on startup when running the prod build tests

* Fix test run

Author: markerikson

__tests__/map-set.js

@@ -367,3 +367,118 @@ function runBaseTest(name, autoFreeze, useListener) {
 		})
 	})
 }
+
+// Bug reproduction: DraftSet leakage when placed in new object and set into Map
+// This bug occurs when a plain object containing a draft Set is set into a DraftMap
+describe("RTK-5159: DraftSet leakage in Map", () => {
+	const {produce} = new Immer({autoFreeze: true})
+
+	test("DraftSet should not leak when placed in new object and set into Map", () => {
+		const baseState = {
+			map: new Map([["key1", {users: new Set()}]])
+		}
+
+		// First produce: add a user
+		const state1 = produce(baseState, draft => {
+			const entry = draft.map.get("key1")
+			entry.users.add({name: "user1"})
+		})
+
+		// Second produce: create new object with existing Set and set into Map
+		const state2 = produce(state1, draft => {
+			const existingUsers = draft.map.get("key1")?.users ?? new Set()
+			const newEntry = {users: existingUsers} // New plain object with existing draft Set
+			draft.map.set("key1", newEntry)
+			newEntry.users.add({name: "user2"})
+		})
+
+		// Should not throw "Cannot use a proxy that has been revoked"
+		const users = state2.map.get("key1").users
+
+		expect(users).toBeInstanceOf(Set)
+		expect([...users].map(u => u.name)).toEqual(["user1", "user2"])
+	})
+
+	test("DraftSet in new object set into Map should finalize correctly", () => {
+		const baseState = {
+			map: new Map([["key1", {items: new Set([1, 2, 3])}]])
+		}
+
+		const result = produce(baseState, draft => {
+			const existingItems = draft.map.get("key1").items
+			// Create a new plain object containing the draft Set
+			const newEntry = {items: existingItems, extra: "data"}
+			draft.map.set("key1", newEntry)
+			newEntry.items.add(4)
+		})
+
+		// Verify the result is properly finalized (not a draft/proxy)
+		const items = result.map.get("key1").items
+		expect(items).toBeInstanceOf(Set)
+		expect([...items]).toEqual([1, 2, 3, 4])
+		expect(result.map.get("key1").extra).toBe("data")
+	})
+
+	// Test for DraftSet.add() with non-draft object containing nested draft
+	// This covers the handleCrossReference() call in DraftSet.add()
+	test("DraftSet.add() should handle non-draft object containing nested draft", () => {
+		const baseState = {
+			items: [{id: 1, name: "item1"}],
+			itemSet: new Set()
+		}
+
+		const result = produce(baseState, draft => {
+			// Get a draft of an existing item
+			const draftItem = draft.items[0]
+			// Create a new plain object that contains the draft
+			const wrapper = {item: draftItem, extra: "wrapper data"}
+			// Add the non-draft wrapper (containing a draft) to the Set
+			draft.itemSet.add(wrapper)
+			// Modify the draft item after adding
+			draftItem.name = "modified"
+		})
+
+		// The wrapper should be finalized correctly
+		const addedItems = [...result.itemSet]
+		expect(addedItems).toHaveLength(1)
+		expect(addedItems[0].item.id).toBe(1)
+		expect(addedItems[0].item.name).toBe("modified")
+		expect(addedItems[0].extra).toBe("wrapper data")
+		// Verify it's not a proxy
+		expect(() => JSON.stringify(result)).not.toThrow()
+	})
+
+	test("DraftSet.add() should handle deeply nested drafts in plain objects", () => {
+		const baseState = {
+			nested: {
+				deep: {
+					value: "original"
+				}
+			},
+			mySet: new Set()
+		}
+
+		const result = produce(baseState, draft => {
+			// Get a draft of a deeply nested object
+			const deepDraft = draft.nested.deep
+			// Create a plain object hierarchy containing the draft
+			const wrapper = {
+				level1: {
+					level2: {
+						draftRef: deepDraft
+					}
+				}
+			}
+			// Add to Set
+			draft.mySet.add(wrapper)
+			// Modify the draft
+			deepDraft.value = "modified"
+		})
+
+		const items = [...result.mySet]
+		expect(items).toHaveLength(1)
+		expect(items[0].level1.level2.draftRef.value).toBe("modified")
+		// Should not throw when accessing
+		expect(() => JSON.stringify(result)).not.toThrow()
+	})
+})

__tests__/patch.js

@@ -1494,3 +1494,42 @@ test("#897 appendPatch", () => {
 		a: [1, 2, 3]
 	})
 })
+
+// Bug reproduction: Patch path truncation in production build
+// This bug is caused by the "key_" in state check failing after minification
+// The patch path should include the full path for nested array modifications
+describe("RTK-5159: Patch path truncation bug", () => {
+	test("patch paths should include full path for nested array modifications", () => {
+		const baseState = {
+			lists: [{items: [{id: 1}]}]
+		}
+
+		const [result, patches] = produceWithPatches(baseState, draft => {
+			draft.lists[0].items.push({id: 2})
+		})
+
+		// The patch path should be ["lists", 0, "items", 1], not just [1]
+		expect(patches[0].path).toEqual(["lists", 0, "items", 1])
+		expect(patches).toEqual([
+			{op: "add", path: ["lists", 0, "items", 1], value: {id: 2}}
+		])
+	})
+
+	test("patch paths correct for deeply nested object in array", () => {
+		const baseState = {
+			queries: {
+				queryKey: {
+					data: {
+						items: [{name: "item1"}]
+					}
+				}
+			}
+		}
+
+		const [result, patches] = produceWithPatches(baseState, draft => {
+			draft.queries.queryKey.data.items.push({name: "item2"})
+		})
+
+		expect(patches[0].path).toEqual(["queries", "queryKey", "data", "items", 1])
+	})
+})

package.json

@@ -26,11 +26,10 @@
   "types": "./dist/immer.d.ts",
   "sideEffects": false,
   "scripts": {
-    "pretest": "yarn build",
     "test": "vitest run && yarn test:build && yarn test:flow",
     "test:perf": "cd __performance_tests__ && node add-data.mjs && node todo.mjs && node incremental.mjs && node large-obj.mjs",
     "test:flow": "yarn flow check __tests__/flow",
-    "test:build": "vitest run --config vitest.config.build.ts",
+    "test:build": "yarn build && vitest run --config vitest.config.build.ts",
     "test:src": "vitest run",
     "watch": "vitest",
     "coverage": "vitest run --coverage",

src/core/finalize.ts

@@ -240,23 +240,29 @@ export function handleCrossReference(
 		target.callbacks_.push(function nestedDraftCleanup() {
 			const targetCopy = latest(target)
 
-			if (get(targetCopy, key, target.type_) === value) {
-				// Process the value to replace any nested drafts
-				// finalizeAssigned(target, key, target.scope_)
-
-				if (
-					scope_.drafts_.length > 1 &&
-					((target as Exclude<ImmerState, SetState>).assigned_!.get(key) ??
-						false) === true &&
-					target.copy_
-				) {
-					// This might be a non-draft value that has drafts
-					// inside. We do need to recurse here to handle those.
-					handleValue(
-						get(target.copy_, key, target.type_),
-						scope_.handledSet_,
-						scope_
-					)
+			// For Sets, check if value is still in the set
+			if (target.type_ === ArchType.Set) {
+				if (targetCopy.has(value)) {
+					// Process the value to replace any nested drafts
+					handleValue(value, scope_.handledSet_, scope_)
+				}
+			} else {
+				// Maps/objects
+				if (get(targetCopy, key, target.type_) === value) {
+					if (
+						scope_.drafts_.length > 1 &&
+						((target as Exclude<ImmerState, SetState>).assigned_!.get(key) ??
+							false) === true &&
+						target.copy_
+					) {
+						// This might be a non-draft value that has drafts
+						// inside. We do need to recurse here to handle those.
+						handleValue(
+							get(target.copy_, key, target.type_),
+							scope_.handledSet_,
+							scope_
+						)
+					}
 				}
 			}
 		})

src/core/proxy.ts

@@ -1,5 +1,4 @@
 import {
-	each,
 	has,
 	is,
 	isDraftable,
@@ -17,7 +16,6 @@ import {
 	die,
 	createProxy,
 	ArchType,
-	ImmerScope,
 	handleCrossReference,
 	WRITABLE,
 	CONFIGURABLE,
@@ -256,14 +254,17 @@ export const objectTraps: ProxyHandler<ProxyState> = {
  */
 
 const arrayTraps: ProxyHandler<[ProxyArrayState]> = {}
-each(objectTraps, (key, fn) => {
+// Use `for..in` instead of `each` to work around a weird
+// prod test suite issue
+for (let key in objectTraps) {
+	let fn = objectTraps[key as keyof typeof objectTraps] as Function
 	// @ts-ignore
 	arrayTraps[key] = function() {
 		const args = arguments
 		args[0] = args[0][0]
 		return fn.apply(this, args)
 	}
-})
+}
 arrayTraps.deleteProperty = function(state, prop) {
 	if (process.env.NODE_ENV !== "production" && isNaN(parseInt(prop as any)))
 		die(13)

src/plugins/mapset.ts

@@ -16,7 +16,8 @@ import {
 	ArchType,
 	each,
 	getValue,
-	PluginMapSet
+	PluginMapSet,
+	handleCrossReference
 } from "../internal"
 
 export function enableMapSet() {
@@ -58,6 +59,7 @@ export function enableMapSet() {
 				state.assigned_!.set(key, true)
 				state.copy_!.set(key, value)
 				state.assigned_!.set(key, true)
+				handleCrossReference(state, key, value)
 			}
 			return this
 		}
@@ -222,6 +224,7 @@ export function enableMapSet() {
 				prepareSetCopy(state)
 				markChanged(state)
 				state.copy_!.add(value)
+				handleCrossReference(state, value, value)
 			}
 			return this
 		}

src/plugins/patches.ts

@@ -49,7 +49,7 @@ export function enablePatches() {
 
 	function getPath(state: ImmerState, path: PatchPath = []): PatchPath | null {
 		// Step 1: Check if state has a stored key
-		if ("key_" in state && state.key_ !== undefined) {
+		if (state.key_ !== undefined) {
 			// Step 2: Validate the key is still valid in parent
 
 			const parentCopy = state.parent_!.copy_ ?? state.parent_!.base_

vitest.config.build.ts

@@ -1,11 +1,38 @@
 import {defineConfig} from "vitest/config"
 import path from "path"
 
+const prodCJSPath = path.resolve(__dirname, "dist/cjs/immer.cjs.production.js")
+
 export default defineConfig({
 	resolve: {
-		alias: {
-			"^src/(.*)": path.resolve(__dirname, "dist/cjs/immer.cjs.production.js")
-		}
+		// 	// Make all `immer` imports use the production build
+		alias: [
+			{
+				find: /^src\/(.*)/,
+				replacement: prodCJSPath
+			},
+			{
+				find: "../src/immer",
+				replacement: prodCJSPath
+			},
+			{
+				find: "immer",
+				replacement: prodCJSPath
+			}
+		],
+		// Ensure only one copy of immer is used throughout the dependency tree
+		dedupe: ["immer"]
+	},
+	// Force Vite to process immer so our alias applies to dependencies too
+	optimizeDeps: {
+		include: ["immer"],
+		// Force re-bundling to pick up our alias
+		force: true
+	},
+	// SSR settings are needed because Vitest runs in Node environment
+	ssr: {
+		// Don't externalize immer - this makes our alias apply to it
+		noExternal: ["immer"]
 	},
 	define: {
 		"global.USES_BUILD": true,
@@ -18,6 +45,15 @@ export default defineConfig({
 		resolveSnapshotPath: (testPath: string, snapExtension: string) =>
 			testPath.replace("__tests__", "__tests__/__prod_snapshots__") +
 			snapExtension,
-		reporters: ["default", "./vitest-custom-reporter.ts"]
+		reporters: ["default", "./vitest-custom-reporter.ts"],
+		// Ensure deps are processed through Vite's transform pipeline
+		deps: {
+			optimizer: {
+				// For SSR (Node) environment, include immer so it's transformed
+				ssr: {
+					include: ["immer"]
+				}
+			}
+		}
 	}
 })