fix: address PR review feedback for task history migration

daniel-lxs · daniel-lxs · commit 020567f16bd1 · 2025-07-02T15:46:09.000-05:00
- 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
diff --git 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<K extends WorkspaceSettingsKey>(key: K, value: WorkspaceSettings[K]) {
-		this.workspaceStateCache[key] = value
-		return this.originalContext.workspaceState.update(key, value)
+	async updateWorkspaceState<K extends WorkspaceSettingsKey>(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
@@ -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
@@ -135,45 +135,72 @@ export async function migrateTaskHistoryWithContextProxy(
 	workspaceFolder: vscode.WorkspaceFolder | undefined,
 ): Promise<void> {
 	try {
-		const alreadyMigrated = context.globalState.get<boolean>(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<boolean>(migrationKey)
 		if (alreadyMigrated) {
 			return
 		}
 
 		if (!workspaceFolder) {
+			// Migration skipped: no workspace folder
 			return
 		}
 		// Get the raw global state directly from context
 		const rawGlobalState = context.globalState.get<any>("globalSettings", {})
 		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)
-
-		// Get tasks that don't belong to current workspace
-		const otherWorkspaceTasks = taskHistory.filter((task) => task.workspace !== currentWorkspacePath)
-
-		if (workspaceTasks.length > 0) {
-			// Get existing workspace settings
-			const workspaceSettings = contextProxy.getWorkspaceSettings()
-			const existingWorkspaceHistory = workspaceSettings.taskHistory || []
-
-			// 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]
-
-			// Update workspace state with the merged history
-			await contextProxy.updateWorkspaceState("taskHistory", mergedHistory)
+		const otherWorkspaceTasks = taskHistory.filter(
+			(task) => task.workspace && task.workspace !== currentWorkspacePath,
+		)
+		const tasksWithoutWorkspace = taskHistory.filter((task) => !task.workspace)
+
+		// 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`,
+		)
+
+		// Migrate tasks for current workspace and tasks without workspace
+		const tasksToMigrate = [...workspaceTasks, ...tasksWithoutWorkspace]
+
+		if (tasksToMigrate.length > 0) {
+			try {
+				// Get existing workspace settings
+				const workspaceSettings = contextProxy.getWorkspaceSettings()
+				const existingWorkspaceHistory = workspaceSettings.taskHistory || []
+
+				// 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}`)
+		}
 	}
 }
