From fdc598f6e6d799f1649c004f25c488d1cf7b9ffa Mon Sep 17 00:00:00 2001
From: littensy <56808540+littensy@users.noreply.github.com>
Date: Wed, 5 Nov 2025 12:35:25 -0800
Subject: [PATCH] fix: allow dispose in observers

---
 .../charm/src/__tests__/observe.spec.luau     | 36 +++++++++++++++++++
 packages/charm/src/observe.luau               | 21 ++++++++---
 2 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/packages/charm/src/__tests__/observe.spec.luau b/packages/charm/src/__tests__/observe.spec.luau
index 3f9a835..f308477 100644
--- a/packages/charm/src/__tests__/observe.spec.luau
+++ b/packages/charm/src/__tests__/observe.spec.luau
@@ -132,4 +132,40 @@ test("should handle self removal", function()
 	assert(table.concat(edits) == "+1+2+3+4-1-2-3-4", `bad edits: {table.concat(edits)}`)
 end)
 
+test("should allow disposal during addition", function()
+	local items, setItems = signal({})
+	local edits = {}
+	local dispose
+
+	dispose = observe(items, function(item)
+		if item == 1 then
+			dispose()
+		end
+		table.insert(edits, `+{item}`)
+		return function()
+			table.insert(edits, `-{item}`)
+		end
+	end)
+
+	setItems({ 1, 2, 3, 4 })
+	assert(table.concat(edits) == "+1-1", `bad edits: {table.concat(edits)}`)
+end)
+
+test("should allow disposal during removal", function()
+	local items, setItems = signal({ 1, 2, 3, 4 })
+	local edits = {}
+	local dispose
+
+	dispose = observe(items, function(item)
+		table.insert(edits, `+{item}`)
+		return function()
+			dispose()
+			table.insert(edits, `-{item}`)
+		end
+	end)
+
+	setItems({ 1, 2, 3, nil :: any, 5 })
+	assert(table.concat(edits) == "+1+2+3+4-1-2-3-4", `bad edits: {table.concat(edits)}`)
+end)
+
 return {}
diff --git a/packages/charm/src/observe.luau b/packages/charm/src/observe.luau
index 8511f43..2a6b9b9 100644
--- a/packages/charm/src/observe.luau
+++ b/packages/charm/src/observe.luau
@@ -19,12 +19,16 @@ local recursive = core.recursive
 --- @see https://github.com/littensy/charm?tab=readme-ov-file#observegetter-callback
 local function observe<K, V>(getter: () -> { [K]: V }, callback: (item: V, key: K) -> ...(() -> ())?): () -> ()
 	local innerScopes: { [K]: () -> () } = {}
+	local disposed = false
 
 	local function updateScopes(data: { [K]: V })
 		for key, dispose in innerScopes do
 			if data[key] == nil then
 				innerScopes[key] = nil
 				dispose()
+				if disposed then
+					return
+				end
 			end
 		end
 
@@ -32,9 +36,15 @@ local function observe<K, V>(getter: () -> { [K]: V }, callback: (item: V, key:
 			if innerScopes[key] == nil then
 				-- Detach the inner scope from the parent effect to prevent
 				-- disposing when the effect re-runs.
-				innerScopes[key] = effectScope(function()
+				local stopScope = effectScope(function()
 					return callback(value, key)
 				end, true)
+
+				if disposed then
+					return stopScope()
+				else
+					innerScopes[key] = stopScope
+				end
 			end
 		end
 	end
@@ -51,10 +61,13 @@ local function observe<K, V>(getter: () -> { [K]: V }, callback: (item: V, key:
 		end)
 
 		return function()
-			for _, dispose in innerScopes do
-				dispose()
+			if not disposed then
+				disposed = true
+				for key, dispose in innerScopes do
+					innerScopes[key] = nil
+					dispose()
+				end
 			end
-			table.clear(innerScopes)
 		end
 	end)
 end