diff --git a/src/core/task-persistence/TaskHistoryStore.ts b/src/core/task-persistence/TaskHistoryStore.ts index 219918b029..c6c3c6910f 100644 --- a/src/core/task-persistence/TaskHistoryStore.ts +++ b/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 { - 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 { + return this.withLock(async () => { + this.cache.clear() + }) } // ────────────────────────────── Migration ────────────────────────────── diff --git a/src/core/task-persistence/__tests__/TaskHistoryStore.spec.ts b/src/core/task-persistence/__tests__/TaskHistoryStore.spec.ts index 75e0a16253..3b7e9041a4 100644 --- a/src/core/task-persistence/__tests__/TaskHistoryStore.spec.ts +++ b/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((resolve) => { + signalWriteStarted = resolve + }) + let releaseWrite!: () => void + const writeCanFinish = new Promise((resolve) => { + releaseWrite = resolve + }) + let releaseStaleRead!: () => void + const staleReadCanFinish = new Promise((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((resolve) => { + signalWriteStarted = resolve + }) + let releaseWrite!: () => void + const writeCanFinish = new Promise((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()", () => { diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 01abc5db98..c3918e24a5 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/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()