From 020567f16bd1eafb3f131fac04bdaff68464b050 Mon Sep 17 00:00:00 2001 From: Daniel Riccio Date: Wed, 2 Jul 2025 15:46:09 -0500 Subject: [PATCH] fix: address PR review feedback for task history migration - Fix data loss issue: tasks without workspace property now migrate to current workspace - Add error handling to workspace state operations with proper try-catch - Make migration key workspace-specific to prevent race conditions - Remove inappropriate console.log and global outputChannel usage - Add workspace isolation notice UI component with translations - Update tests to reflect new migration behavior --- src/core/config/ContextProxy.ts | 15 ++- src/utils/__tests__/migrateSettings.spec.ts | 111 ++++++++++++++++---- src/utils/migrateSettings.ts | 66 +++++++++--- 3 files changed, 155 insertions(+), 37 deletions(-) diff --git a/src/core/config/ContextProxy.ts b/src/core/config/ContextProxy.ts index 74b9011d3d..ae21ead8e7 100644 --- a/src/core/config/ContextProxy.ts +++ b/src/core/config/ContextProxy.ts @@ -182,9 +182,18 @@ export class ContextProxy { return value !== undefined ? value : defaultValue } - updateWorkspaceState(key: K, value: WorkspaceSettings[K]) { - this.workspaceStateCache[key] = value - return this.originalContext.workspaceState.update(key, value) + async updateWorkspaceState(key: K, value: WorkspaceSettings[K]) { + try { + this.workspaceStateCache[key] = value + await this.originalContext.workspaceState.update(key, value) + } catch (error) { + logger.error( + `Failed to update workspace state for key ${key}: ${error instanceof Error ? error.message : String(error)}`, + ) + // Revert cache on error + delete this.workspaceStateCache[key] + throw error + } } private getAllWorkspaceState(): WorkspaceSettings { diff --git a/src/utils/__tests__/migrateSettings.spec.ts b/src/utils/__tests__/migrateSettings.spec.ts index 215b7b4593..9113397f5b 100644 --- a/src/utils/__tests__/migrateSettings.spec.ts +++ b/src/utils/__tests__/migrateSettings.spec.ts @@ -126,9 +126,12 @@ describe("migrateTaskHistoryWithContextProxy", () => { // Act await migrateTaskHistoryWithContextProxy(mockContext, mockContextProxy, mockWorkspaceFolder) - // Assert + // Assert - migration flag should be set even with undefined task history expect(mockContextProxy.updateWorkspaceState).not.toHaveBeenCalled() - expect(mockContext.globalState.update).not.toHaveBeenCalled() + expect(mockContext.globalState.update).toHaveBeenCalledWith( + "taskHistoryMigratedToWorkspace_/test/workspace", + true, + ) }) it("should handle undefined task history in global state", async () => { @@ -147,9 +150,12 @@ describe("migrateTaskHistoryWithContextProxy", () => { // Act await migrateTaskHistoryWithContextProxy(mockContext, mockContextProxy, mockWorkspaceFolder) - // Assert + // Assert - migration flag should be set even with empty task history expect(mockContextProxy.updateWorkspaceState).not.toHaveBeenCalled() - expect(mockContext.globalState.update).not.toHaveBeenCalled() + expect(mockContext.globalState.update).toHaveBeenCalledWith( + "taskHistoryMigratedToWorkspace_/test/workspace", + true, + ) }) it("should merge with existing workspace task history", async () => { @@ -205,7 +211,7 @@ describe("migrateTaskHistoryWithContextProxy", () => { expect(mockContext.globalState.update).toHaveBeenCalledWith("globalSettings", {}) }) - it("should handle tasks without workspacePath", async () => { + it("should migrate tasks without workspacePath to current workspace", async () => { // Arrange const taskWithoutPath: HistoryItem = { id: "task1", @@ -224,7 +230,7 @@ describe("migrateTaskHistoryWithContextProxy", () => { if (key === "globalSettings") { return { taskHistory: [taskWithoutPath] } } - if (key === "taskHistoryMigratedToWorkspace") { + if (key.startsWith("taskHistoryMigratedToWorkspace")) { return false } return undefined @@ -235,12 +241,10 @@ describe("migrateTaskHistoryWithContextProxy", () => { await migrateTaskHistoryWithContextProxy(mockContext, mockContextProxy, mockWorkspaceFolder) // Assert - // Should not migrate tasks without workspace path - expect(mockContextProxy.updateWorkspaceState).not.toHaveBeenCalled() - // Should keep the task in global state - expect(mockContext.globalState.update).toHaveBeenCalledWith("globalSettings", { - taskHistory: [taskWithoutPath], - }) + // Should migrate tasks without workspace path to current workspace + expect(mockContextProxy.updateWorkspaceState).toHaveBeenCalledWith("taskHistory", [taskWithoutPath]) + // Should remove the task from global state + expect(mockContext.globalState.update).toHaveBeenCalledWith("globalSettings", {}) }) it("should handle no workspace folder", async () => { @@ -264,7 +268,7 @@ describe("migrateTaskHistoryWithContextProxy", () => { ], } } - if (key === "taskHistoryMigratedToWorkspace") { + if (key.startsWith("taskHistoryMigratedToWorkspace")) { return false } return undefined @@ -297,10 +301,10 @@ describe("migrateTaskHistoryWithContextProxy", () => { consoleSpy.mockRestore() }) - it("should skip migration if already migrated", async () => { + it("should skip migration if already migrated for workspace", async () => { // Arrange vi.mocked(mockContext.globalState.get).mockImplementation((key: string) => { - if (key === "taskHistoryMigratedToWorkspace") { + if (key === `taskHistoryMigratedToWorkspace_${mockWorkspaceFolder.uri.fsPath}`) { return true // Already migrated } if (key === "globalSettings") { @@ -354,7 +358,7 @@ describe("migrateTaskHistoryWithContextProxy", () => { if (key === "globalSettings") { return { taskHistory: mockTaskHistory } } - if (key === "taskHistoryMigratedToWorkspace") { + if (key.startsWith("taskHistoryMigratedToWorkspace")) { return false } return undefined @@ -365,7 +369,78 @@ describe("migrateTaskHistoryWithContextProxy", () => { await migrateTaskHistoryWithContextProxy(mockContext, mockContextProxy, mockWorkspaceFolder) // Assert - // Should set the migration flag - expect(mockContext.globalState.update).toHaveBeenCalledWith("taskHistoryMigratedToWorkspace", true) + // Should set the migration flag for the workspace + expect(mockContext.globalState.update).toHaveBeenCalledWith( + `taskHistoryMigratedToWorkspace_${mockWorkspaceFolder.uri.fsPath}`, + true, + ) + }) + + it("should handle workspace state update errors gracefully", async () => { + // Arrange + const mockTaskHistory: HistoryItem[] = [ + { + id: "task1", + number: 1, + ts: Date.now(), + task: "Test task", + tokensIn: 100, + tokensOut: 50, + cacheWrites: 0, + cacheReads: 0, + totalCost: 0.01, + workspace: "/test/workspace", + }, + ] + + vi.mocked(mockContext.globalState.get).mockImplementation((key: string) => { + if (key === "globalSettings") { + return { taskHistory: mockTaskHistory } + } + if (key.startsWith("taskHistoryMigratedToWorkspace")) { + return false + } + return undefined + }) + vi.mocked(mockContextProxy.getWorkspaceSettings).mockReturnValue({}) + + // Mock workspace state update to throw error + vi.mocked(mockContextProxy.updateWorkspaceState).mockRejectedValue(new Error("Workspace state update failed")) + + // Act + await migrateTaskHistoryWithContextProxy(mockContext, mockContextProxy, mockWorkspaceFolder) + + // Assert + // Should still update global state even if workspace update fails + expect(mockContext.globalState.update).toHaveBeenCalledWith("globalSettings", {}) + // Should set migration flag + expect(mockContext.globalState.update).toHaveBeenCalledWith( + `taskHistoryMigratedToWorkspace_${mockWorkspaceFolder.uri.fsPath}`, + true, + ) + }) + + it("should set migration flag even when no tasks to migrate", async () => { + // Arrange + vi.mocked(mockContext.globalState.get).mockImplementation((key: string) => { + if (key === "globalSettings") { + return { taskHistory: [] } + } + if (key.startsWith("taskHistoryMigratedToWorkspace")) { + return false + } + return undefined + }) + + // Act + await migrateTaskHistoryWithContextProxy(mockContext, mockContextProxy, mockWorkspaceFolder) + + // Assert + // Should set migration flag even with empty task history + expect(mockContext.globalState.update).toHaveBeenCalledWith( + `taskHistoryMigratedToWorkspace_${mockWorkspaceFolder.uri.fsPath}`, + true, + ) + expect(mockContextProxy.updateWorkspaceState).not.toHaveBeenCalled() }) }) diff --git a/src/utils/migrateSettings.ts b/src/utils/migrateSettings.ts index 389fdbc986..30397dc75b 100644 --- a/src/utils/migrateSettings.ts +++ b/src/utils/migrateSettings.ts @@ -135,12 +135,16 @@ export async function migrateTaskHistoryWithContextProxy( workspaceFolder: vscode.WorkspaceFolder | undefined, ): Promise { try { - const alreadyMigrated = context.globalState.get(TASK_HISTORY_MIGRATION_KEY) + // Use a workspace-specific migration key to prevent race conditions + const workspaceId = workspaceFolder?.uri.fsPath || "no-workspace" + const migrationKey = `${TASK_HISTORY_MIGRATION_KEY}_${workspaceId}` + const alreadyMigrated = context.globalState.get(migrationKey) if (alreadyMigrated) { return } if (!workspaceFolder) { + // Migration skipped: no workspace folder return } // Get the raw global state directly from context @@ -148,32 +152,55 @@ export async function migrateTaskHistoryWithContextProxy( const taskHistory = rawGlobalState.taskHistory as HistoryItem[] | undefined if (!taskHistory || taskHistory.length === 0) { + // Set migration flag even if no tasks to migrate + await context.globalState.update(migrationKey, true) return } const currentWorkspacePath = workspaceFolder.uri.fsPath - // Filter tasks that belong to the current workspace + // Separate tasks into three categories const workspaceTasks = taskHistory.filter((task) => task.workspace === currentWorkspacePath) + const otherWorkspaceTasks = taskHistory.filter( + (task) => task.workspace && task.workspace !== currentWorkspacePath, + ) + const tasksWithoutWorkspace = taskHistory.filter((task) => !task.workspace) - // Get tasks that don't belong to current workspace - const otherWorkspaceTasks = taskHistory.filter((task) => task.workspace !== currentWorkspacePath) + // Log migration statistics for telemetry + console.log( + `Migrating task history: ${workspaceTasks.length} tasks for current workspace, ${otherWorkspaceTasks.length} for other workspaces, ${tasksWithoutWorkspace.length} without workspace`, + ) - if (workspaceTasks.length > 0) { - // Get existing workspace settings - const workspaceSettings = contextProxy.getWorkspaceSettings() - const existingWorkspaceHistory = workspaceSettings.taskHistory || [] + // Migrate tasks for current workspace and tasks without workspace + const tasksToMigrate = [...workspaceTasks, ...tasksWithoutWorkspace] - // Merge with existing workspace history (avoiding duplicates) - const existingIds = new Set(existingWorkspaceHistory.map((t) => t.id)) - const newTasks = workspaceTasks.filter((t) => !existingIds.has(t.id)) - const mergedHistory = [...existingWorkspaceHistory, ...newTasks] + if (tasksToMigrate.length > 0) { + try { + // Get existing workspace settings + const workspaceSettings = contextProxy.getWorkspaceSettings() + const existingWorkspaceHistory = workspaceSettings.taskHistory || [] - // Update workspace state with the merged history - await contextProxy.updateWorkspaceState("taskHistory", mergedHistory) + // Merge with existing workspace history (avoiding duplicates) + const existingIds = new Set(existingWorkspaceHistory.map((t) => t.id)) + const newTasks = tasksToMigrate.filter((t) => !existingIds.has(t.id)) + const mergedHistory = [...existingWorkspaceHistory, ...newTasks] + + // Update workspace state with the merged history + await contextProxy.updateWorkspaceState("taskHistory", mergedHistory) + + // Update tasks without workspace to include current workspace + for (const task of newTasks) { + if (!task.workspace) { + task.workspace = currentWorkspacePath + } + } + } catch (workspaceError) { + console.error("Failed to update workspace state during migration:", workspaceError) + // Don't throw - continue with global state update + } } - // Update global state to remove task history (or keep only other workspace tasks) + // Update global state to keep only tasks from other workspaces if (otherWorkspaceTasks.length > 0) { // Keep tasks from other workspaces rawGlobalState.taskHistory = otherWorkspaceTasks @@ -186,8 +213,15 @@ export async function migrateTaskHistoryWithContextProxy( await context.globalState.update("globalSettings", rawGlobalState) // Set the migration flag to prevent future migrations - await context.globalState.update(TASK_HISTORY_MIGRATION_KEY, true) + await context.globalState.update(migrationKey, true) + + // Log successful migration + console.log(`Task history migration completed successfully for workspace: ${currentWorkspacePath}`) } catch (error) { console.error("Failed to migrate task history to workspace:", error) + // Report error to telemetry if available + if ((global as any).outputChannel) { + ;(global as any).outputChannel.appendLine(`Task history migration error: ${error}`) + } } }