diff --git a/src/core/task-persistence/__tests__/taskMessages.spec.ts b/src/core/task-persistence/__tests__/taskMessages.spec.ts index e494b5b594..bb0fb4aee4 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" let tmpBaseDir: string @@ -101,6 +101,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 89ec372c61..5ca53cfe2d 100644 --- a/src/core/task-persistence/taskMessages.ts +++ b/src/core/task-persistence/taskMessages.ts @@ -65,6 +65,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 874e2f8f08..e4175f6245 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -99,8 +99,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" @@ -128,7 +129,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 { REQUESTY_BASE_URL } from "../../shared/utils/requesty" @@ -2298,12 +2299,130 @@ 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 + } + + 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 + } + + const updatedHistoryItem = { ...historyItem, workspace: currentWorkspace } + await this.resetTaskCheckpointsForWorkspaceChange(historyItem, updatedHistoryItem) + return updatedHistoryItem + } + + private async resetTaskCheckpointsForWorkspaceChange( + originalHistoryItem: HistoryItem, + updatedHistoryItem: HistoryItem, + ): Promise { + const taskId = originalHistoryItem.id + const globalStoragePath = this.contextProxy.globalStorageUri.fsPath + let removedCheckpoints: ClineMessage[] = [] + 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)) + }, + }) + await this.updateTaskHistory(updatedHistoryItem) + } 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 { + if (this.taskHistoryStore.get(taskId)?.workspace === updatedHistoryItem.workspace) { + await this.updateTaskHistory(originalHistoryItem) + } + } 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) + } + } + } + } + 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..4daf9e7b62 --- /dev/null +++ b/src/core/webview/__tests__/ClineProvider.history-workspace.spec.ts @@ -0,0 +1,392 @@ +import fs from "fs/promises" +import os from "os" +import path from "path" +import * as vscode from "vscode" + +import type { ClineMessage, HistoryItem } from "@roo-code/types" + +import { ClineProvider } from "../ClineProvider" +import * as taskMessages from "../../task-persistence/taskMessages" + +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() + }) + + it.each([true, false])("resets checkpoint rows with an existing checkpoint directory: %s", async (hasDirectory) => { + const storagePath = await fs.mkdtemp(path.join(os.tmpdir(), "zoo-history-workspace-")) + const taskDir = path.join(storagePath, "tasks", "task-1602") + const checkpointsDir = path.join(taskDir, "checkpoints") + await fs.mkdir(taskDir, { recursive: true }) + if (hasDirectory) { + await fs.mkdir(checkpointsDir) + await fs.writeFile(path.join(checkpointsDir, "HEAD"), "old checkpoint") + } + await fs.writeFile( + path.join(taskDir, "ui_messages.json"), + 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 = createProvider("/current/workspace") + Object.defineProperty(provider, "contextProxy", { + value: { globalStorageUri: { fsPath: storagePath } }, + }) + const original = historyItem("/old/worktree") + const updated = { ...original, workspace: "/current/workspace" } + provider.updateTaskHistory = vi.fn().mockResolvedValue([]) + Object.defineProperty(provider, "taskHistoryStore", { value: { get: vi.fn().mockReturnValue(original) } }) + + try { + await provider["resetTaskCheckpointsForWorkspaceChange"](original, updated) + + await expect(fs.stat(checkpointsDir)).rejects.toMatchObject({ code: "ENOENT" }) + const messages = JSON.parse(await fs.readFile(path.join(taskDir, "ui_messages.json"), "utf8")) + expect(messages).toMatchObject([ + { type: "say", say: "task", text: "Continue" }, + { type: "say", say: "text", text: "Still useful" }, + ]) + expect(provider.updateTaskHistory).toHaveBeenCalledWith(updated) + } finally { + await fs.rm(storagePath, { recursive: true, force: true }) + } + }) + + it.each([false, true])("preserves intervening message writes (rollback: %s)", async (rollback) => { + const storagePath = await fs.mkdtemp(path.join(os.tmpdir(), "zoo-history-concurrent-")) + const taskDir = path.join(storagePath, "tasks", "task-1602") + const checkpointsDir = path.join(taskDir, "checkpoints") + await fs.mkdir(checkpointsDir, { recursive: true }) + const original = historyItem("/old/worktree") + const updated = { ...original, workspace: "/current/workspace" } + const options = { taskId: original.id, globalStoragePath: storagePath } + const checkpoint: ClineMessage = { ts: 2, type: "say", say: "checkpoint_saved", text: "old-hash" } + await taskMessages.saveTaskMessages({ + ...options, + messages: [{ ts: 1, type: "say", say: "text", text: "before" }, checkpoint], + }) + const provider = createProvider(updated.workspace) + Object.defineProperty(provider, "contextProxy", { value: { globalStorageUri: { fsPath: storagePath } } }) + Object.defineProperty(provider, "taskHistoryStore", { value: { get: () => original } }) + 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" }], + }) + } + }) + provider.updateTaskHistory = vi.fn().mockImplementation(async () => { + 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 [updated] + }) + try { + 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" }, + ]) + } finally { + await fs.rm(storagePath, { recursive: true, force: true }) + } + }) + + it("restores checkpoint files and messages when workspace persistence fails", async () => { + const storagePath = await fs.mkdtemp(path.join(os.tmpdir(), "zoo-history-workspace-rollback-")) + const taskDir = path.join(storagePath, "tasks", "task-1602") + const checkpointsDir = path.join(taskDir, "checkpoints") + const messagesPath = path.join(taskDir, "ui_messages.json") + const originalMessages = [ + { type: "say", say: "task", ts: 1, text: "Continue" }, + { type: "say", say: "checkpoint_saved", ts: 2, text: "old-hash" }, + ] + await fs.mkdir(checkpointsDir, { recursive: true }) + await fs.writeFile(path.join(checkpointsDir, "HEAD"), "old checkpoint") + await fs.writeFile(messagesPath, JSON.stringify(originalMessages)) + + const provider = createProvider("/current/workspace") + const original = historyItem("/old/worktree") + const updated = { ...original, workspace: "/current/workspace" } + Object.defineProperty(provider, "contextProxy", { + value: { globalStorageUri: { fsPath: storagePath } }, + }) + Object.defineProperty(provider, "taskHistoryStore", { value: { get: vi.fn().mockReturnValue(original) } }) + provider.updateTaskHistory = vi.fn().mockRejectedValue(new Error("history write failed")) + + try { + 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) + } finally { + await fs.rm(storagePath, { recursive: true, force: true }) + } + }) + + 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 }) => { + const storagePath = await fs.mkdtemp(path.join(os.tmpdir(), "zoo-history-workspace-rollback-failure-")) + const taskDir = path.join(storagePath, "tasks", "task-1602") + const checkpointsDir = path.join(taskDir, "checkpoints") + const messagesPath = path.join(taskDir, "ui_messages.json") + const originalMessages = [ + { type: "say", say: "task", ts: 1, text: "Continue" }, + { type: "say", say: "checkpoint_saved", ts: 2, text: "old-hash" }, + ] + await fs.mkdir(checkpointsDir, { recursive: true }) + await fs.writeFile(path.join(checkpointsDir, "HEAD"), "old checkpoint") + await fs.writeFile(messagesPath, JSON.stringify(originalMessages)) + + const provider = createProvider("/current/workspace") + const original = historyItem("/old/worktree") + const updated = { ...original, workspace: "/current/workspace" } + Object.defineProperty(provider, "contextProxy", { + value: { globalStorageUri: { fsPath: storagePath } }, + }) + // History can already be committed when broadcasting the update fails. + Object.defineProperty(provider, "taskHistoryStore", { value: { get: vi.fn().mockReturnValue(updated) } }) + const originalError = new Error("history broadcast failed") + const updateHistory = vi.fn().mockRejectedValueOnce(originalError) + if (failures.includes("history")) { + updateHistory.mockRejectedValueOnce(new Error("history restore failed")) + } else { + updateHistory.mockResolvedValueOnce([original]) + } + provider.updateTaskHistory = updateHistory + provider["log"] = vi.fn() + 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) + }) + + try { + 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(updateHistory).toHaveBeenCalledTimes(2) + expect(updateHistory).toHaveBeenLastCalledWith(original) + 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") + } + } finally { + await fs.rm(storagePath, { recursive: true, force: true }) + } + }, + ) + + it("retries checkpoint backup cleanup without rolling back committed history", async () => { + const storagePath = await fs.mkdtemp(path.join(os.tmpdir(), "zoo-history-workspace-cleanup-")) + const taskDir = path.join(storagePath, "tasks", "task-1602") + const checkpointsDir = path.join(taskDir, "checkpoints") + const messagesPath = path.join(taskDir, "ui_messages.json") + await fs.mkdir(checkpointsDir, { recursive: true }) + 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" }, + ]), + ) + + const provider = createProvider("/current/workspace") + const original = historyItem("/old/worktree") + const updated = { ...original, workspace: "/current/workspace" } + Object.defineProperty(provider, "contextProxy", { + value: { globalStorageUri: { fsPath: storagePath } }, + }) + Object.defineProperty(provider, "taskHistoryStore", { value: { get: vi.fn().mockReturnValue(updated) } }) + provider.updateTaskHistory = vi.fn().mockResolvedValue([]) + provider["log"] = vi.fn() + vi.spyOn(fs, "rm").mockRejectedValueOnce(new Error("backup cleanup failed")) + + try { + await expect(provider["resetTaskCheckpointsForWorkspaceChange"](original, updated)).resolves.toBeUndefined() + + expect(provider.updateTaskHistory).toHaveBeenCalledOnce() + expect(provider.updateTaskHistory).toHaveBeenCalledWith(updated) + 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() + await fs.rm(storagePath, { recursive: true, force: true }) + } + }) +}) 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..3dac9d35d0 --- /dev/null +++ b/src/extension/__tests__/api-resume-task.spec.ts @@ -0,0 +1,41 @@ +import * as vscode from "vscode" + +import type { HistoryItem } from "@roo-code/types" + +import { API } from "../api" +import type { ClineProvider } from "../../core/webview/ClineProvider" + +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) + + 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 316e7a6c9d..5bc81b4329 100644 --- a/src/extension/api.ts +++ b/src/extension/api.ts @@ -218,7 +218,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" })