diff --git a/apps/vscode-e2e/src/visual/electron.visual.ts b/apps/vscode-e2e/src/visual/electron.visual.ts index 8d4072ea7b..28ff0f4a36 100644 --- a/apps/vscode-e2e/src/visual/electron.visual.ts +++ b/apps/vscode-e2e/src/visual/electron.visual.ts @@ -122,6 +122,8 @@ async function startScene( mock.addFixture({ match: { userMessage: CHAT_PROMPT }, response: { + // Keep the context meter independent of system prompts and temporary workspace paths. + usage: { prompt_tokens: 4096, completion_tokens: 128, total_tokens: 4224 }, toolCalls: [ { name: "attempt_completion", @@ -230,6 +232,9 @@ for (const scenario of scenarios) { await contentFrame .getByText(scenario.landmark, { exact: false }) .waitFor({ state: "visible", timeout: 60_000 }) + if (scenario.scene === "chat") { + await expect(contentFrame.getByTestId("context-tokens-count")).toHaveText("4.2k") + } await contentFrame.evaluate(() => { if (document.activeElement instanceof HTMLElement) document.activeElement.blur() }) diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 86ce5d8e67..5a44e4e415 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" @@ -2305,12 +2306,103 @@ 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 + const messages = await readTaskMessages({ taskId, globalStoragePath }) + const messagesWithoutCheckpoints = messages.filter( + (message) => !(message.type === "say" && message.say === "checkpoint_saved"), + ) + 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 { + if (messagesWithoutCheckpoints.length !== messages.length) { + await saveTaskMessages({ messages: messagesWithoutCheckpoints, taskId, globalStoragePath }) + } + await this.updateTaskHistory(updatedHistoryItem) + } catch (error) { + if (messagesWithoutCheckpoints.length !== messages.length) { + await saveTaskMessages({ messages, taskId, globalStoragePath }) + } + if (checkpointDirectoryStaged) { + await fs.rename(checkpointBackupDir, checkpointsDir) + } + if (this.taskHistoryStore.get(taskId)?.workspace === updatedHistoryItem.workspace) { + await this.updateTaskHistory(originalHistoryItem) + } + throw error + } + + if (checkpointDirectoryStaged) { + try { + await fs.rm(checkpointBackupDir, { recursive: true, force: true }) + } catch (error) { + this.log( + `[resetTaskCheckpointsForWorkspaceChange] Failed to remove checkpoint backup for ${taskId}: ${error instanceof Error ? error.message : String(error)}`, + ) + } + } + } + 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..b84e52de73 --- /dev/null +++ b/src/core/webview/__tests__/ClineProvider.history-workspace.spec.ts @@ -0,0 +1,239 @@ +import fs from "fs/promises" +import os from "os" +import path from "path" +import * as vscode from "vscode" + +import type { HistoryItem } from "@roo-code/types" + +import { ClineProvider } from "../ClineProvider" + +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("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("removes the old checkpoint repository and checkpoint-only chat rows", async () => { + 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(checkpointsDir, { recursive: true }) + 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("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("keeps committed history when checkpoint backup cleanup fails", 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")) + 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..8c4b5043ef --- /dev/null +++ b/src/extension/__tests__/api-resume-task.spec.ts @@ -0,0 +1,33 @@ +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", + task: "Resume task", + tokensIn: 0, + tokensOut: 0, + totalCost: 0, +} as HistoryItem + +describe("API.resumeTask", () => { + it("does not restore a task when workspace selection is cancelled", async () => { + const provider = { + viewLaunched: true, + context: {}, + getTaskWithId: vi.fn().mockResolvedValue({ historyItem }), + prepareHistoryItemForResume: vi.fn().mockResolvedValue(undefined), + createTaskWithHistoryItem: vi.fn(), + on: vi.fn(), + } as unknown as ClineProvider + const outputChannel = { appendLine: vi.fn() } as unknown as vscode.OutputChannel + const api = new API(outputChannel, provider) + + await api.resumeTask(historyItem.id) + + 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" })