diff --git a/src/core/task-persistence/__tests__/taskMessages.spec.ts b/src/core/task-persistence/__tests__/taskMessages.spec.ts index 6b968fe96c..95f4602950 100644 --- a/src/core/task-persistence/__tests__/taskMessages.spec.ts +++ b/src/core/task-persistence/__tests__/taskMessages.spec.ts @@ -19,7 +19,7 @@ vi.mock("../../../utils/safeWriteJson", () => ({ })) // Import after mocks -import { saveTaskMessages, readTaskMessages } from "../taskMessages" +import { saveTaskMessages, readTaskMessages, updateTaskMessages } from "../taskMessages" import { formatErrorDetail } from "../readFileWithMissingRetry" let tmpBaseDir: string @@ -119,6 +119,55 @@ describe("taskMessages.saveTaskMessages", () => { }) }) +describe("taskMessages.updateTaskMessages", () => { + it("keeps legacy same-timestamp identities distinct across removal and rollback", async () => { + const { safeWriteJson } = + await vi.importActual("../../../utils/safeWriteJson") + const options = { taskId: "task-legacy-update", globalStoragePath: tmpBaseDir } + hoisted.safeWriteJsonMock.mockImplementationOnce(safeWriteJson) + await saveTaskMessages({ ...options, messages: [] }) + await fs.writeFile( + path.join(tmpBaseDir, "tasks", options.taskId, "ui_messages.json"), + JSON.stringify([ + { ts: 1, type: "say", say: "checkpoint_saved", text: "checkpoint" }, + { ts: 1, type: "say", say: "text", text: "keep" }, + ]), + ) + let removed: ClineMessage[] = [] + hoisted.safeWriteJsonMock.mockImplementationOnce(safeWriteJson) + await updateTaskMessages({ + ...options, + update: (messages) => { + removed = messages.slice(0, 1) + return messages.slice(1) + }, + }) + hoisted.safeWriteJsonMock.mockImplementationOnce(safeWriteJson) + await saveTaskMessages({ ...options, messages: removed, merge: true }) + const messages = await readTaskMessages(options) + expect(messages.map((message) => message.text).sort()).toEqual(["checkpoint", "keep"]) + expect(new Set(messages.map((message) => message.messageId)).size).toBe(2) + }) + + it.each([null, "invalid", {}])("rejects a missing or invalid locked snapshot: %j", async (existing: unknown) => { + const { safeWriteJson } = + await vi.importActual("../../../utils/safeWriteJson") + hoisted.safeWriteJsonMock.mockImplementationOnce(safeWriteJson) + const taskId = "task-invalid-update" + const taskDir = path.join(tmpBaseDir, "tasks", taskId) + await fs.mkdir(taskDir, { recursive: true }) + const filePath = path.join(taskDir, "ui_messages.json") + if (existing !== null) await fs.writeFile(filePath, JSON.stringify(existing)) + const update = vi.fn((messages: ClineMessage[]) => messages) + await expect(updateTaskMessages({ taskId, globalStoragePath: tmpBaseDir, update })).rejects.toMatchObject({ + kind: "invalid", + }) + expect(update).not.toHaveBeenCalled() + if (existing !== null) expect(JSON.parse(await fs.readFile(filePath, "utf8"))).toEqual(existing) + else await expect(fs.stat(filePath)).rejects.toMatchObject({ code: "ENOENT" }) + }) +}) + describe("taskMessages.readTaskMessages", () => { it("rejects invalid JSON without treating it as empty history", async () => { const taskId = "task-corrupt-json" diff --git a/src/core/task-persistence/taskMessages.ts b/src/core/task-persistence/taskMessages.ts index b67ad2b7f8..961c1790f0 100644 --- a/src/core/task-persistence/taskMessages.ts +++ b/src/core/task-persistence/taskMessages.ts @@ -73,6 +73,24 @@ export type SaveTaskMessagesOptions = { merge?: boolean } +/** Update the current disk snapshot under the same lock used by every message writer. */ +export async function updateTaskMessages({ + taskId, + globalStoragePath, + update, +}: ReadTaskMessagesOptions & { update: (messages: ClineMessage[]) => ClineMessage[] }): Promise { + const taskDir = await getTaskDirectoryPath(globalStoragePath, taskId) + const filePath = path.join(taskDir, GlobalFileNames.uiMessages) + await safeWriteJson(filePath, null, { + merge: (existing) => { + if (!Array.isArray(existing)) { + throw new TaskMessagesReadError("invalid", `Cannot update missing or invalid messages for ${taskId}`) + } + return ensureMessageIdentifiers(update(ensureMessageIdentifiers(existing))) + }, + }) +} + export async function saveTaskMessages({ messages, taskId, diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index bc8d366dd3..917979b32a 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -101,8 +101,9 @@ import { SkillsManager } from "../../services/skills/SkillsManager" import { fileExistsAtPath } from "../../utils/fs" import { setTtsEnabled, setTtsSpeed } from "../../utils/tts" import { getWorkspaceGitInfo } from "../../utils/git" -import { getWorkspacePath } from "../../utils/path" +import { arePathsEqual, getWorkspacePath } from "../../utils/path" import { OrganizationAllowListViolationError } from "../../utils/errors" +import { getTaskDirectoryPath } from "../../utils/storage" import { setPanel } from "../../activate/registerCommands" @@ -131,7 +132,7 @@ import { interruptDelegatedChild, LifecycleTransitionError, } from "../task-persistence" -import { readTaskMessages } from "../task-persistence/taskMessages" +import { readTaskMessages, updateTaskMessages } from "../task-persistence/taskMessages" import { getNonce } from "./getNonce" import { getUri } from "./getUri" import { omitOriginalContentFromExtensionMessage } from "./stripOriginalContent" @@ -2430,12 +2431,184 @@ export class ClineProvider if (id !== this.getCurrentTask()?.taskId) { // Non-current task. const { historyItem } = await this.getTaskWithId(id) - await this.createTaskWithHistoryItem(historyItem) // Clears existing task. + const preparedHistoryItem = await this.prepareHistoryItemForResume(historyItem) + if (!preparedHistoryItem) { + return + } + await this.createTaskWithHistoryItem(preparedHistoryItem) // Clears existing task. } await this.postMessageToWebview({ type: "action", action: "chatButtonClicked" }) } + public async prepareHistoryItemForResume(historyItem: T): Promise { + const currentWorkspace = this.cwd + const originalWorkspace = historyItem.workspace + + if (!currentWorkspace || !originalWorkspace || arePathsEqual(currentWorkspace, originalWorkspace)) { + return historyItem + } + + // A live Task owns this history and would re-add checkpoint rows or rewrite its workspace after a reset. + if (this.isTaskLiveInAnyProvider(historyItem.id)) { + return historyItem + } + + const useCurrentWorkspace = { title: "Use Current Workspace" } + const openOriginalWorkspace = { title: "Open Original Workspace" } + const selection = await vscode.window.showWarningMessage( + `This conversation was created in "${originalWorkspace}", but the current workspace is "${currentWorkspace}". ` + + "Choose where to continue. Using the current workspace resets checkpoints created in the original workspace.", + { modal: true }, + useCurrentWorkspace, + openOriginalWorkspace, + ) + + if (selection?.title === openOriginalWorkspace.title) { + await vscode.commands.executeCommand("vscode.openFolder", vscode.Uri.file(originalWorkspace), { + forceNewWindow: true, + }) + return undefined + } + + if (selection?.title !== useCurrentWorkspace.title) { + return undefined + } + + // The modal has no timeout, so the task may have been opened elsewhere while it was shown. + if (this.isTaskLiveInAnyProvider(historyItem.id)) { + this.log(`[prepareHistoryItemForResume] Task ${historyItem.id} became live during workspace selection`) + return undefined + } + + const updatedHistoryItem = { ...historyItem, workspace: currentWorkspace } + const persisted = await this.resetTaskCheckpointsForWorkspaceChange(historyItem, updatedHistoryItem) + // Resume from the record as persisted now, not the snapshot read before the prompt. + return { ...historyItem, ...persisted, workspace: currentWorkspace } + } + + private isTaskLiveInAnyProvider(taskId: string): boolean { + return ClineProvider.getAllInstances().some((provider) => provider.taskRegistry.getById(taskId) !== undefined) + } + + /** + * Persist only `workspace`, read-modify-writing the stored record so fields that changed meanwhile + * (status, pendingAction, ...) are never overwritten from a stale snapshot. With `expectedWorkspace`, + * the stored record is left unchanged unless it still has that workspace. + */ + private async setTaskHistoryWorkspace( + historyItem: HistoryItem, + workspace: string | undefined, + expectedWorkspace?: string, + ): Promise { + const taskId = historyItem.id + if (!this.taskHistoryStore.get(taskId)) { + // Legacy fallback before the store is initialized: there is no stored record to preserve. + if (expectedWorkspace !== undefined) return undefined + await this.updateTaskHistory({ ...historyItem, workspace }) + return this.taskHistoryStore.get(taskId) + } + + let changed = false + await this.taskHistoryStore.atomicReadAndUpdate(taskId, (current) => { + if (expectedWorkspace !== undefined && current.workspace !== expectedWorkspace) return current + changed = current.workspace !== workspace + return { ...current, workspace } + }) + this.recentTasksCache = undefined + const updated = this.taskHistoryStore.get(taskId) + if (changed && updated && this.isViewLaunched) { + await this.postMessageToWebview({ type: "taskHistoryItemUpdated", taskHistoryItem: updated }) + } + return updated + } + + private async resetTaskCheckpointsForWorkspaceChange( + originalHistoryItem: HistoryItem, + updatedHistoryItem: HistoryItem, + ): Promise { + const taskId = originalHistoryItem.id + const globalStoragePath = this.contextProxy.globalStorageUri.fsPath + let removedCheckpoints: ClineMessage[] = [] + let persisted: HistoryItem | undefined + const taskDir = await getTaskDirectoryPath(globalStoragePath, taskId) + const checkpointsDir = path.join(taskDir, "checkpoints") + const checkpointBackupDir = path.join(taskDir, `checkpoints.workspace-change-${crypto.randomUUID()}`) + let checkpointDirectoryStaged = false + + try { + await fs.rename(checkpointsDir, checkpointBackupDir) + checkpointDirectoryStaged = true + } catch (error) { + if (!(error instanceof Error && "code" in error && error.code === "ENOENT")) { + throw error + } + } + + try { + await updateTaskMessages({ + taskId, + globalStoragePath, + update: (messages) => { + removedCheckpoints = messages.filter( + (message) => message.type === "say" && message.say === "checkpoint_saved", + ) + return messages.filter((message) => !removedCheckpoints.includes(message)) + }, + }) + persisted = await this.setTaskHistoryWorkspace(originalHistoryItem, updatedHistoryItem.workspace) + } catch (error) { + if (removedCheckpoints.length > 0) { + try { + // Restore only the removed rows; never replace intervening message writes. + await saveTaskMessages({ messages: removedCheckpoints, taskId, globalStoragePath, merge: true }) + } catch (rollbackError) { + this.log( + `[resetTaskCheckpointsForWorkspaceChange] Failed to restore messages for ${taskId}: ${rollbackError instanceof Error ? rollbackError.message : String(rollbackError)}`, + ) + } + } + if (checkpointDirectoryStaged) { + try { + await fs.rename(checkpointBackupDir, checkpointsDir) + } catch (rollbackError) { + this.log( + `[resetTaskCheckpointsForWorkspaceChange] Failed to restore checkpoints for ${taskId} from ${checkpointBackupDir}: ${rollbackError instanceof Error ? rollbackError.message : String(rollbackError)}`, + ) + } + } + try { + // Revert only the workspace, and only if no newer write has moved it again. + await this.setTaskHistoryWorkspace( + originalHistoryItem, + originalHistoryItem.workspace, + updatedHistoryItem.workspace, + ) + } catch (rollbackError) { + this.log( + `[resetTaskCheckpointsForWorkspaceChange] Failed to restore history for ${taskId}: ${rollbackError instanceof Error ? rollbackError.message : String(rollbackError)}`, + ) + } + throw error + } + + if (checkpointDirectoryStaged) { + for (let attempt = 0; attempt < 3; attempt++) { + try { + await fs.rm(checkpointBackupDir, { recursive: true, force: true }) + break + } catch (error) { + this.log( + `[resetTaskCheckpointsForWorkspaceChange] Failed to remove checkpoint backup ${checkpointBackupDir}: ${error instanceof Error ? error.message : String(error)}`, + ) + if (attempt < 2) await delay(100) + } + } + } + + return persisted + } + async exportTaskWithId(id: string) { const { historyItem, apiConversationHistory } = await this.getTaskWithId(id) const fileName = getTaskFileName(historyItem.ts) diff --git a/src/core/webview/__tests__/ClineProvider.history-workspace.spec.ts b/src/core/webview/__tests__/ClineProvider.history-workspace.spec.ts new file mode 100644 index 0000000000..2343fc3035 --- /dev/null +++ b/src/core/webview/__tests__/ClineProvider.history-workspace.spec.ts @@ -0,0 +1,463 @@ +import fs from "fs/promises" +import os from "os" +import path from "path" +import * as vscode from "vscode" + +import type { HistoryItem, PendingTaskAction } from "@roo-code/types" + +import { ClineProvider } from "../ClineProvider" +import type { Task } from "../../task/Task" +import { TaskRegistry } from "../../task/TaskRegistry" +import { TaskHistoryStore } from "../../task-persistence/TaskHistoryStore" +import * as taskMessages from "../../task-persistence/taskMessages" +import { GlobalFileNames } from "../../../shared/globalFileNames" + +const historyItem = (workspace: string): HistoryItem => ({ + id: "task-1602", + number: 1, + ts: 1, + task: "Continue in another worktree", + tokensIn: 0, + tokensOut: 0, + cacheWrites: 0, + cacheReads: 0, + totalCost: 0, + workspace, +}) + +const createProvider = (workspace: string) => { + const provider = Object.create(ClineProvider.prototype) as ClineProvider + Object.defineProperty(provider, "currentWorkspacePath", { value: workspace, writable: true }) + return provider +} + +describe("ClineProvider historical workspace selection", () => { + afterEach(() => { + vi.restoreAllMocks() + }) + + it("keeps a history item unchanged when it already belongs to the current workspace", async () => { + const provider = createProvider("/current/workspace") + const prompt = vi.spyOn(vscode.window, "showWarningMessage") + const item = historyItem("/current/workspace") + + await expect(provider.prepareHistoryItemForResume(item)).resolves.toBe(item) + expect(prompt).not.toHaveBeenCalled() + }) + + it.each([ + { current: "", original: "/old/worktree" }, + { current: "/current/workspace", original: undefined }, + ])("does not prompt when a workspace is absent ($current, $original)", async ({ current, original }) => { + const provider = createProvider(current) + const prompt = vi.spyOn(vscode.window, "showWarningMessage") + const item = { ...historyItem("/old/worktree"), workspace: original } + await expect(provider.prepareHistoryItemForResume(item)).resolves.toBe(item) + expect(prompt).not.toHaveBeenCalled() + }) + + it("cancels restoration without mutating history when the mismatch prompt is dismissed", async () => { + const provider = createProvider("/current/workspace") + vi.spyOn(vscode.window, "showWarningMessage").mockResolvedValue(undefined) + provider["resetTaskCheckpointsForWorkspaceChange"] = vi.fn() + provider.updateTaskHistory = vi.fn() + + await expect(provider.prepareHistoryItemForResume(historyItem("/old/worktree"))).resolves.toBeUndefined() + expect(provider["resetTaskCheckpointsForWorkspaceChange"]).not.toHaveBeenCalled() + expect(provider.updateTaskHistory).not.toHaveBeenCalled() + }) + + it("opens the original workspace in a new window without restoring the task", async () => { + const provider = createProvider("/current/workspace") + const prompt = vi + .spyOn(vscode.window, "showWarningMessage") + .mockResolvedValue({ title: "Open Original Workspace" }) + const executeCommand = vi.spyOn(vscode.commands, "executeCommand").mockResolvedValue(undefined) + + await expect(provider.prepareHistoryItemForResume(historyItem("/old/worktree"))).resolves.toBeUndefined() + expect(prompt).toHaveBeenCalledWith( + expect.any(String), + { modal: true }, + { title: "Use Current Workspace" }, + { title: "Open Original Workspace" }, + ) + expect(executeCommand).toHaveBeenCalledWith( + "vscode.openFolder", + expect.objectContaining({ fsPath: "/old/worktree" }), + { forceNewWindow: true }, + ) + }) + + it("moves the task to the current workspace and resets workspace-specific checkpoints", async () => { + const provider = createProvider("/current/workspace") + vi.spyOn(vscode.window, "showWarningMessage").mockResolvedValue({ title: "Use Current Workspace" }) + provider["resetTaskCheckpointsForWorkspaceChange"] = vi.fn().mockResolvedValue(undefined) + provider.updateTaskHistory = vi.fn().mockResolvedValue([]) + const item = historyItem("/old/worktree") + + await expect(provider.prepareHistoryItemForResume(item)).resolves.toEqual({ + ...item, + workspace: "/current/workspace", + }) + expect(provider["resetTaskCheckpointsForWorkspaceChange"]).toHaveBeenCalledWith(item, { + ...item, + workspace: "/current/workspace", + }) + }) + + it("restores history with the workspace selected by the mismatch prompt", async () => { + const provider = createProvider("/current/workspace") + const original = historyItem("/old/worktree") + const prepared = { ...original, workspace: "/current/workspace" } + provider.getCurrentTask = vi.fn().mockReturnValue(undefined) + provider.getTaskWithId = vi.fn().mockResolvedValue({ historyItem: original }) + provider.prepareHistoryItemForResume = vi.fn().mockResolvedValue(prepared) + provider.createTaskWithHistoryItem = vi.fn().mockResolvedValue({}) + provider.postMessageToWebview = vi.fn().mockResolvedValue(true) + + await provider.showTaskWithId(original.id) + + expect(provider.createTaskWithHistoryItem).toHaveBeenCalledWith(prepared) + expect(provider.postMessageToWebview).toHaveBeenCalledWith({ + type: "action", + action: "chatButtonClicked", + }) + }) + + it("does not restore or reveal a task when workspace selection is cancelled", async () => { + const provider = createProvider("/current/workspace") + const original = historyItem("/old/worktree") + provider.getCurrentTask = vi.fn().mockReturnValue(undefined) + provider.getTaskWithId = vi.fn().mockResolvedValue({ historyItem: original }) + provider.prepareHistoryItemForResume = vi.fn().mockResolvedValue(undefined) + provider.createTaskWithHistoryItem = vi.fn() + provider.postMessageToWebview = vi.fn() + + await provider.showTaskWithId(original.id) + + expect(provider.createTaskWithHistoryItem).not.toHaveBeenCalled() + expect(provider.postMessageToWebview).not.toHaveBeenCalled() + }) + + describe("with a persisted task history store", () => { + let storagePath: string + let taskDir: string + let checkpointsDir: string + let messagesPath: string + let store: TaskHistoryStore + + const original = historyItem("/old/worktree") + const updated = { ...original, workspace: "/current/workspace" } + + beforeEach(async () => { + storagePath = await fs.mkdtemp(path.join(os.tmpdir(), "zoo-history-workspace-")) + taskDir = path.join(storagePath, "tasks", original.id) + checkpointsDir = path.join(taskDir, "checkpoints") + messagesPath = path.join(taskDir, "ui_messages.json") + store = new TaskHistoryStore(storagePath) + await store.initialize() + store.dispose() // Keep watcher reconciliation out of these explicit write tests. + await store.upsert(original) + }) + + afterEach(async () => { + ClineProvider["activeInstances"].clear() + await fs.rm(storagePath, { recursive: true, force: true }) + }) + + const originalMessages = [ + { type: "say", say: "task", ts: 1, text: "Continue" }, + { type: "say", say: "checkpoint_saved", ts: 2, text: "old-hash" }, + ] + + const createStoredProvider = () => { + const provider = createProvider("/current/workspace") + Object.defineProperty(provider, "contextProxy", { value: { globalStorageUri: { fsPath: storagePath } } }) + Object.defineProperty(provider, "taskHistoryStore", { value: store }) + provider["log"] = vi.fn() + return provider + } + + const writeCheckpointState = async () => { + await fs.mkdir(checkpointsDir, { recursive: true }) + await fs.writeFile(path.join(checkpointsDir, "HEAD"), "old checkpoint") + await fs.writeFile(messagesPath, JSON.stringify(originalMessages)) + } + + const readPersistedHistory = async () => + JSON.parse(await fs.readFile(path.join(taskDir, GlobalFileNames.historyItem), "utf8")) as HistoryItem + + it.each([true, false])( + "resets checkpoint rows with an existing checkpoint directory: %s", + async (hasDirectory) => { + if (hasDirectory) { + await fs.mkdir(checkpointsDir) + await fs.writeFile(path.join(checkpointsDir, "HEAD"), "old checkpoint") + } + await fs.writeFile( + messagesPath, + JSON.stringify([ + { type: "say", say: "task", ts: 1, text: "Continue" }, + { type: "say", say: "checkpoint_saved", ts: 2, text: "old-hash" }, + { type: "say", say: "text", ts: 3, text: "Still useful" }, + ]), + ) + const provider = createStoredProvider() + + await expect(provider["resetTaskCheckpointsForWorkspaceChange"](original, updated)).resolves.toEqual( + updated, + ) + + await expect(fs.stat(checkpointsDir)).rejects.toMatchObject({ code: "ENOENT" }) + expect(JSON.parse(await fs.readFile(messagesPath, "utf8"))).toMatchObject([ + { type: "say", say: "task", text: "Continue" }, + { type: "say", say: "text", text: "Still useful" }, + ]) + expect(store.get(original.id)?.workspace).toBe("/current/workspace") + expect((await readPersistedHistory()).workspace).toBe("/current/workspace") + }, + ) + + it.each([false, true])( + "keeps history fields written while the prompt was open (rollback: %s)", + async (rollback) => { + await writeCheckpointState() + const pendingAction: PendingTaskAction = { + kind: "create_subtask", + actionId: "action-1", + approvalText: "Create a subtask", + mode: "code", + message: "Do the work", + todos: [], + } + // An explicit stale status is what a full-item write would put back over the newer one. + const snapshot: HistoryItem = { ...original, status: "active" } + await store.upsert(snapshot) + const fresh = { status: "delegated", awaitingChildId: "child-1", pendingAction } as const + const provider = createStoredProvider() + // The stale snapshot is the record read before the modal; another writer updates it meanwhile. + vi.spyOn(vscode.window, "showWarningMessage").mockImplementation(async () => { + await store.atomicReadAndUpdate(original.id, (current) => ({ ...current, ...fresh })) + return { title: "Use Current Workspace" } + }) + const historyBroadcastFailure = new Error("history broadcast failed") + provider.isViewLaunched = rollback + provider.postMessageToWebview = vi.fn().mockRejectedValueOnce(historyBroadcastFailure) + + const prepared = provider.prepareHistoryItemForResume(snapshot) + + const expected = { ...snapshot, ...fresh, workspace: rollback ? original.workspace : updated.workspace } + if (rollback) await expect(prepared).rejects.toBe(historyBroadcastFailure) + else await expect(prepared).resolves.toEqual(expected) + expect(store.get(original.id)).toEqual(expected) + expect(await readPersistedHistory()).toEqual(expected) + }, + ) + + it.each([false, true])("preserves intervening message writes (rollback: %s)", async (rollback) => { + await fs.mkdir(checkpointsDir, { recursive: true }) + const options = { taskId: original.id, globalStoragePath: storagePath } + await taskMessages.saveTaskMessages({ + ...options, + messages: [ + { ts: 1, type: "say", say: "text", text: "before" }, + { ts: 2, type: "say", say: "checkpoint_saved", text: "old-hash" }, + ], + }) + const provider = createStoredProvider() + const rename = fs.rename + vi.spyOn(fs, "rename").mockImplementation(async (source, destination) => { + await rename(source, destination) + if (source === checkpointsDir) { + // A live task saves after reset starts, before the locked removal. + await taskMessages.saveTaskMessages({ + ...options, + merge: true, + messages: [{ ts: 3, type: "say", say: "text", text: "during reset" }], + }) + } + }) + const realAtomicUpdate = store.atomicReadAndUpdate.bind(store) + vi.spyOn(store, "atomicReadAndUpdate").mockImplementationOnce(async (...args) => { + await taskMessages.saveTaskMessages({ + ...options, + merge: true, + messages: [{ ts: 1, type: "say", say: "text", text: "updated during history write" }], + }) + if (rollback) throw new Error("history write failed") + return realAtomicUpdate(...args) + }) + + const reset = provider["resetTaskCheckpointsForWorkspaceChange"](original, updated) + if (rollback) await expect(reset).rejects.toThrow("history write failed") + else await reset + + const messages = await taskMessages.readTaskMessages(options) + expect(messages.map(({ ts, text }) => ({ ts, text }))).toEqual([ + { ts: 1, text: "updated during history write" }, + ...(rollback ? [{ ts: 2, text: "old-hash" }] : []), + { ts: 3, text: "during reset" }, + ]) + }) + + it("restores checkpoint files and messages when workspace persistence fails", async () => { + await writeCheckpointState() + const provider = createStoredProvider() + vi.spyOn(store, "atomicReadAndUpdate").mockRejectedValueOnce(new Error("history write failed")) + + await expect(provider["resetTaskCheckpointsForWorkspaceChange"](original, updated)).rejects.toThrow( + "history write failed", + ) + + await expect(fs.readFile(path.join(checkpointsDir, "HEAD"), "utf8")).resolves.toBe("old checkpoint") + expect(JSON.parse(await fs.readFile(messagesPath, "utf8"))).toMatchObject(originalMessages) + expect(store.get(original.id)?.workspace).toBe("/old/worktree") + }) + + it.each([ + { failures: ["messages"] }, + { failures: ["checkpoints"] }, + { failures: ["history"] }, + { failures: ["messages", "checkpoints", "history"] }, + ])( + "attempts every rollback and preserves the original error when $failures restoration fails", + async ({ failures }) => { + await writeCheckpointState() + const provider = createStoredProvider() + // History is already committed when broadcasting the update fails. + const originalError = new Error("history broadcast failed") + provider.isViewLaunched = true + provider.postMessageToWebview = vi.fn().mockRejectedValueOnce(originalError).mockResolvedValue(true) + const realAtomicUpdate = store.atomicReadAndUpdate.bind(store) + const atomicUpdate = vi.spyOn(store, "atomicReadAndUpdate") + if (failures.includes("history")) { + atomicUpdate + .mockImplementationOnce(realAtomicUpdate) + .mockRejectedValueOnce(new Error("history restore failed")) + } + const save = vi.spyOn(taskMessages, "saveTaskMessages") + if (failures.includes("messages")) { + save.mockRejectedValueOnce(new Error("messages restore failed")) + } + const rename = fs.rename + const renameSpy = vi.spyOn(fs, "rename").mockImplementation(async (source, destination) => { + if (destination === checkpointsDir && failures.includes("checkpoints")) { + throw new Error("checkpoints restore failed") + } + await rename(source, destination) + }) + + await expect(provider["resetTaskCheckpointsForWorkspaceChange"](original, updated)).rejects.toBe( + originalError, + ) + + expect(save).toHaveBeenCalledOnce() + expect(save.mock.calls[0][0]).toMatchObject({ + messages: [originalMessages[1]], + taskId: original.id, + globalStoragePath: storagePath, + merge: true, + }) + expect(renameSpy).toHaveBeenCalledWith( + expect.stringContaining("checkpoints.workspace-change-"), + checkpointsDir, + ) + expect(atomicUpdate).toHaveBeenCalledTimes(2) + expect(store.get(original.id)?.workspace).toBe( + failures.includes("history") ? "/current/workspace" : "/old/worktree", + ) + expect(provider["log"]).toHaveBeenCalledTimes(failures.length) + for (const failure of failures) { + expect(provider["log"]).toHaveBeenCalledWith(expect.stringContaining(`${failure} restore failed`)) + } + if (!failures.includes("messages")) { + expect(JSON.parse(await fs.readFile(messagesPath, "utf8"))).toMatchObject(originalMessages) + } + if (!failures.includes("checkpoints")) { + await expect(fs.readFile(path.join(checkpointsDir, "HEAD"), "utf8")).resolves.toBe("old checkpoint") + } + }, + ) + + it("does not roll back a workspace that a newer write already moved", async () => { + await writeCheckpointState() + const provider = createStoredProvider() + const broadcastFailure = new Error("history broadcast failed") + provider.isViewLaunched = true + provider.postMessageToWebview = vi.fn().mockImplementationOnce(async () => { + await store.atomicReadAndUpdate(original.id, (current) => ({ + ...current, + workspace: "/newer/workspace", + })) + throw broadcastFailure + }) + + await expect(provider["resetTaskCheckpointsForWorkspaceChange"](original, updated)).rejects.toBe( + broadcastFailure, + ) + + expect(store.get(original.id)?.workspace).toBe("/newer/workspace") + }) + + it("retries checkpoint backup cleanup without rolling back committed history", async () => { + await writeCheckpointState() + const provider = createStoredProvider() + vi.spyOn(fs, "rm").mockRejectedValueOnce(new Error("backup cleanup failed")) + + try { + await expect(provider["resetTaskCheckpointsForWorkspaceChange"](original, updated)).resolves.toEqual( + updated, + ) + + expect(store.get(original.id)?.workspace).toBe("/current/workspace") + expect(provider["log"]).toHaveBeenCalledWith(expect.stringContaining("backup cleanup failed")) + expect(vi.mocked(fs.rm)).toHaveBeenCalledTimes(2) + expect( + (await fs.readdir(taskDir)).filter((name) => name.startsWith("checkpoints.workspace-change-")), + ).toEqual([]) + await expect(fs.stat(checkpointsDir)).rejects.toMatchObject({ code: "ENOENT" }) + expect(JSON.parse(await fs.readFile(messagesPath, "utf8"))).toMatchObject([ + { type: "say", say: "task", text: "Continue" }, + ]) + } finally { + vi.mocked(fs.rm).mockRestore() + } + }) + + it("leaves a task owned by a live Task untouched without prompting", async () => { + await writeCheckpointState() + const provider = createStoredProvider() + const registry = new TaskRegistry() + // TaskRegistry reads only taskId; a full Task would start unrelated services. + registry.push({ taskId: original.id } as Task) + Object.defineProperty(provider, "taskRegistry", { value: registry }) + ClineProvider["activeInstances"].add(provider) + const prompt = vi.spyOn(vscode.window, "showWarningMessage") + + await expect(provider.prepareHistoryItemForResume(original)).resolves.toBe(original) + + expect(prompt).not.toHaveBeenCalled() + expect(store.get(original.id)?.workspace).toBe("/old/worktree") + await expect(fs.readFile(path.join(checkpointsDir, "HEAD"), "utf8")).resolves.toBe("old checkpoint") + }) + + it("cancels without resetting when the task becomes live while the prompt is open", async () => { + await writeCheckpointState() + const provider = createStoredProvider() + const registry = new TaskRegistry() + Object.defineProperty(provider, "taskRegistry", { value: registry }) + ClineProvider["activeInstances"].add(provider) + vi.spyOn(vscode.window, "showWarningMessage").mockImplementation(async () => { + // TaskRegistry reads only taskId; a full Task would start unrelated services. + registry.push({ taskId: original.id } as Task) + return { title: "Use Current Workspace" } + }) + + await expect(provider.prepareHistoryItemForResume(original)).resolves.toBeUndefined() + + expect(store.get(original.id)?.workspace).toBe("/old/worktree") + await expect(fs.readFile(path.join(checkpointsDir, "HEAD"), "utf8")).resolves.toBe("old checkpoint") + expect(JSON.parse(await fs.readFile(messagesPath, "utf8"))).toMatchObject(originalMessages) + }) + }) +}) diff --git a/src/extension/__tests__/api-resume-task.spec.ts b/src/extension/__tests__/api-resume-task.spec.ts new file mode 100644 index 0000000000..8385f319d5 --- /dev/null +++ b/src/extension/__tests__/api-resume-task.spec.ts @@ -0,0 +1,42 @@ +import * as vscode from "vscode" + +import type { HistoryItem } from "@roo-code/types" + +import { API } from "../api" +import type { ClineProvider } from "../../core/webview/ClineProvider" +import { makeClineProviderFactory } from "../../test-utils/provider" + +const historyItem = { + id: "task-1602", + number: 1, + ts: 1, + task: "Resume task", + tokensIn: 0, + tokensOut: 0, + totalCost: 0, +} satisfies HistoryItem + +describe("API.resumeTask", () => { + it.each([false, true])("resumes only with the prepared history item (accepted: %s)", async (accepted) => { + const prepared = accepted ? { ...historyItem, workspace: "/selected/workspace" } : undefined + // API uses only this provider surface; constructing a full provider starts unrelated services. + const provider = { + viewLaunched: true, + context: {}, + getTaskWithId: vi.fn().mockResolvedValue({ historyItem }), + prepareHistoryItemForResume: vi.fn().mockResolvedValue(prepared), + createTaskWithHistoryItem: vi.fn(), + postMessageToWebview: vi.fn().mockResolvedValue(true), + on: vi.fn(), + } as unknown as ClineProvider + // Only appendLine is used by this API test. + const outputChannel = { appendLine: vi.fn() } as unknown as vscode.OutputChannel + const api = new API(outputChannel, provider, makeClineProviderFactory()) + + await api.resumeTask(historyItem.id) + + expect(provider.prepareHistoryItemForResume).toHaveBeenCalledWith(historyItem) + if (accepted) expect(provider.createTaskWithHistoryItem).toHaveBeenCalledExactlyOnceWith(prepared) + else expect(provider.createTaskWithHistoryItem).not.toHaveBeenCalled() + }) +}) diff --git a/src/extension/api.ts b/src/extension/api.ts index 3a33510346..2234d914e8 100644 --- a/src/extension/api.ts +++ b/src/extension/api.ts @@ -266,7 +266,11 @@ export class API extends EventEmitter implements RooCodeAPI { await this.waitForWebviewLaunch(5_000) const { historyItem } = await this.sidebarProvider.getTaskWithId(taskId) - await this.sidebarProvider.createTaskWithHistoryItem(historyItem) + const preparedHistoryItem = await this.sidebarProvider.prepareHistoryItemForResume(historyItem) + if (!preparedHistoryItem) { + return + } + await this.sidebarProvider.createTaskWithHistoryItem(preparedHistoryItem) if (this.sidebarProvider.viewLaunched) { await this.sidebarProvider.postMessageToWebview({ type: "action", action: "chatButtonClicked" })