fix(task-history): route invalidate() and invalidateAll() through withLock (#912)

* fix: lock task history invalidation

* fix(task-history): address invalidation review feedback

Author: morgan-coded

src/core/task-persistence/TaskHistoryStore.ts

@@ -430,23 +430,27 @@ export class TaskHistoryStore {
 	 * Invalidate a single task's cache entry (re-read from disk on next access).
 	 */
 	async invalidate(taskId: string): Promise<void> {
-		try {
-			const item = await this.readTaskFile(taskId)
-			if (item) {
-				this.cache.set(taskId, item)
-			} else {
+		return this.withLock(async () => {
+			try {
+				const item = await this.readTaskFile(taskId)
+				if (item) {
+					this.cache.set(taskId, item)
+				} else {
+					this.cache.delete(taskId)
+				}
+			} catch {
 				this.cache.delete(taskId)
 			}
-		} catch {
-			this.cache.delete(taskId)
-		}
+		})
 	}
 
 	/**
-	 * Clear all in-memory cache and reload from index.
+	 * Clear all in-memory cache entries; a subsequent `reconcile()` repopulates them from task files.
 	 */
-	invalidateAll(): void {
-		this.cache.clear()
+	async invalidateAll(): Promise<void> {
+		return this.withLock(async () => {
+			this.cache.clear()
+		})
 	}
 
 	// ────────────────────────────── Migration ──────────────────────────────

src/core/task-persistence/__tests__/TaskHistoryStore.spec.ts

@@ -440,6 +440,107 @@ describe("TaskHistoryStore", () => {
 
 			expect(store.get("gone-task")).toBeUndefined()
 		})
+
+		it("waits for an in-flight write before refreshing the cache", async () => {
+			await store.initialize()
+
+			const item = makeHistoryItem({ id: "invalidate-locked", tokensIn: 100 })
+			await store.upsert(item)
+
+			let signalWriteStarted!: () => void
+			const writeStarted = new Promise<void>((resolve) => {
+				signalWriteStarted = resolve
+			})
+			let releaseWrite!: () => void
+			const writeCanFinish = new Promise<void>((resolve) => {
+				releaseWrite = resolve
+			})
+			let releaseStaleRead!: () => void
+			const staleReadCanFinish = new Promise<void>((resolve) => {
+				releaseStaleRead = resolve
+			})
+			let writeReleased = false
+
+			const storeAny = store as any
+			const originalWriteTaskFile = storeAny.writeTaskFile.bind(store)
+			const originalReadTaskFile = storeAny.readTaskFile.bind(store)
+			vi.spyOn(storeAny, "writeTaskFile").mockImplementation(async (...args: unknown[]) => {
+				const next = args[0] as HistoryItem
+				if (next.id === item.id && next.tokensIn === 999) {
+					signalWriteStarted()
+					await writeCanFinish
+				}
+				return originalWriteTaskFile(...args)
+			})
+			vi.spyOn(storeAny, "readTaskFile").mockImplementation(async (...args: unknown[]) => {
+				if (args[0] === item.id && !writeReleased) {
+					await staleReadCanFinish
+					return item
+				}
+				return originalReadTaskFile(...args)
+			})
+
+			const write = store.upsert({ ...item, tokensIn: 999 })
+			await writeStarted
+			const invalidation = store.invalidate(item.id)
+
+			writeReleased = true
+			releaseWrite()
+			await write
+			releaseStaleRead()
+			await invalidation
+
+			expect(store.get(item.id)?.tokensIn).toBe(999)
+		})
+	})
+
+	describe("invalidateAll()", () => {
+		it("waits for an in-flight write before clearing the cache", async () => {
+			const onWrite = vi.fn().mockResolvedValue(undefined)
+			store = new TaskHistoryStore(tmpDir, { onWrite })
+			await store.initialize()
+
+			const first = makeHistoryItem({ id: "invalidate-all-first", ts: 1000, tokensIn: 100 })
+			const second = makeHistoryItem({ id: "invalidate-all-second", ts: 2000 })
+			await store.upsert(first)
+			await store.upsert(second)
+			onWrite.mockClear()
+
+			let signalWriteStarted!: () => void
+			const writeStarted = new Promise<void>((resolve) => {
+				signalWriteStarted = resolve
+			})
+			let releaseWrite!: () => void
+			const writeCanFinish = new Promise<void>((resolve) => {
+				releaseWrite = resolve
+			})
+
+			const storeAny = store as any
+			const originalWriteTaskFile = storeAny.writeTaskFile.bind(store)
+			vi.spyOn(storeAny, "writeTaskFile").mockImplementation(async (...args: unknown[]) => {
+				const item = args[0] as HistoryItem
+				if (item.id === first.id && item.tokensIn === 999) {
+					signalWriteStarted()
+					await writeCanFinish
+				}
+				return originalWriteTaskFile(...args)
+			})
+
+			const write = store.upsert({ ...first, tokensIn: 999 })
+			await writeStarted
+			const invalidation = store.invalidateAll()
+
+			releaseWrite()
+			await write
+			await invalidation
+
+			expect(onWrite).toHaveBeenCalledTimes(1)
+			expect(onWrite.mock.calls[0][0].map((item: HistoryItem) => item.id).sort()).toEqual([
+				"invalidate-all-first",
+				"invalidate-all-second",
+			])
+			expect(store.getAll()).toEqual([])
+		})
 	})
 
 	describe("atomicUpdatePair()", () => {

src/core/webview/webviewMessageHandler.ts

@@ -963,7 +963,7 @@ export const webviewMessageHandler = async (
 
 				// Refresh history whenever Roo tasks were found — even if all already existed —
 				// so a retry after a partial-copy failure still reconciles the store.
-				provider.taskHistoryStore.invalidateAll()
+				await provider.taskHistoryStore.invalidateAll()
 				await provider.taskHistoryStore.reconcile()
 				await provider.taskHistoryStore.flushIndex()
 				await provider.postStateToWebview()