fix: eliminate UI flicker during task cancellation (#9037)

daniel-lxs · web-flow · commit e25b4e62d1e5 · 2025-11-04T21:40:39.000-05:00
* fix: eliminate UI flicker during task cancellation

- Modify ClineProvider.createTaskWithHistoryItem() to detect when rehydrating current task
- Implement in-place task replacement to avoid empty stack state that causes UI flicker
- Add comprehensive unit tests for flicker-free cancel behavior
- Maintain backward compatibility and proper event listener cleanup

Fixes the jarring navigation to home view when cancelling tasks

* fix: ensure proper garbage collection of old task during flicker-free rehydration

- Call abortTask(true) on old task before replacement to stop processes and mark as abandoned
- This ensures proper cleanup and prevents memory leaks during task cancellation
- Add test verification for abortTask cleanup
diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts
@@ -858,7 +858,13 @@ export class ClineProvider
 	}
 
 	public async createTaskWithHistoryItem(historyItem: HistoryItem & { rootTask?: Task; parentTask?: Task }) {
-		await this.removeClineFromStack()
+		// Check if we're rehydrating the current task to avoid flicker
+		const currentTask = this.getCurrentTask()
+		const isRehydratingCurrentTask = currentTask && currentTask.taskId === historyItem.id
+
+		if (!isRehydratingCurrentTask) {
+			await this.removeClineFromStack()
+		}
 
 		// If the history item has a saved mode, restore it and its associated API configuration.
 		if (historyItem.mode) {
@@ -932,11 +938,46 @@ export class ClineProvider
 			enableBridge: BridgeOrchestrator.isEnabled(cloudUserInfo, taskSyncEnabled),
 		})
 
-		await this.addClineToStack(task)
+		if (isRehydratingCurrentTask) {
+			// Replace the current task in-place to avoid UI flicker
+			const stackIndex = this.clineStack.length - 1
 
-		this.log(
-			`[createTaskWithHistoryItem] ${task.parentTask ? "child" : "parent"} task ${task.taskId}.${task.instanceId} instantiated`,
-		)
+			// Properly dispose of the old task to ensure garbage collection
+			const oldTask = this.clineStack[stackIndex]
+
+			// Abort the old task to stop running processes and mark as abandoned
+			try {
+				await oldTask.abortTask(true)
+			} catch (e) {
+				this.log(
+					`[createTaskWithHistoryItem] abortTask() failed for old task ${oldTask.taskId}.${oldTask.instanceId}: ${e.message}`,
+				)
+			}
+
+			// Remove event listeners from the old task
+			const cleanupFunctions = this.taskEventListeners.get(oldTask)
+			if (cleanupFunctions) {
+				cleanupFunctions.forEach((cleanup) => cleanup())
+				this.taskEventListeners.delete(oldTask)
+			}
+
+			// Replace the task in the stack
+			this.clineStack[stackIndex] = task
+			task.emit(RooCodeEventName.TaskFocused)
+
+			// Perform preparation tasks and set up event listeners
+			await this.performPreparationTasks(task)
+
+			this.log(
+				`[createTaskWithHistoryItem] rehydrated task ${task.taskId}.${task.instanceId} in-place (flicker-free)`,
+			)
+		} else {
+			await this.addClineToStack(task)
+
+			this.log(
+				`[createTaskWithHistoryItem] ${task.parentTask ? "child" : "parent"} task ${task.taskId}.${task.instanceId} instantiated`,
+			)
+		}
 
 		// Check if there's a pending edit after checkpoint restoration
 		const operationId = `task-${task.taskId}`
diff --git a/src/core/webview/__tests__/ClineProvider.flicker-free-cancel.spec.ts b/src/core/webview/__tests__/ClineProvider.flicker-free-cancel.spec.ts
@@ -0,0 +1,321 @@
+import { beforeEach, describe, expect, it, vi } from "vitest"
+import * as vscode from "vscode"
+
+import { ClineProvider } from "../ClineProvider"
+import { Task } from "../../task/Task"
+import { ContextProxy } from "../../config/ContextProxy"
+import type { ProviderSettings, HistoryItem } from "@roo-code/types"
+
+// Mock dependencies
+vi.mock("vscode", () => {
+	const mockDisposable = { dispose: vi.fn() }
+	return {
+		workspace: {
+			getConfiguration: vi.fn(() => ({
+				get: vi.fn().mockReturnValue([]),
+				update: vi.fn().mockResolvedValue(undefined),
+			})),
+			workspaceFolders: [],
+			onDidChangeConfiguration: vi.fn(() => mockDisposable),
+		},
+		env: {
+			uriScheme: "vscode",
+			language: "en",
+		},
+		EventEmitter: vi.fn().mockImplementation(() => ({
+			event: vi.fn(),
+			fire: vi.fn(),
+		})),
+		Disposable: {
+			from: vi.fn(),
+		},
+		window: {
+			showErrorMessage: vi.fn(),
+			createTextEditorDecorationType: vi.fn().mockReturnValue({
+				dispose: vi.fn(),
+			}),
+			onDidChangeActiveTextEditor: vi.fn(() => mockDisposable),
+		},
+		Uri: {
+			file: vi.fn().mockReturnValue({ toString: () => "file://test" }),
+		},
+	}
+})
+
+vi.mock("../../task/Task")
+vi.mock("../../config/ContextProxy")
+vi.mock("../../../services/mcp/McpServerManager", () => ({
+	McpServerManager: {
+		getInstance: vi.fn().mockResolvedValue({
+			registerClient: vi.fn(),
+		}),
+		unregisterProvider: vi.fn(),
+	},
+}))
+vi.mock("../../../services/marketplace")
+vi.mock("../../../integrations/workspace/WorkspaceTracker")
+vi.mock("../../config/ProviderSettingsManager")
+vi.mock("../../config/CustomModesManager")
+vi.mock("../../../utils/path", () => ({
+	getWorkspacePath: vi.fn().mockReturnValue("/test/workspace"),
+}))
+
+// Mock TelemetryService
+vi.mock("@roo-code/telemetry", () => ({
+	TelemetryService: {
+		instance: {
+			setProvider: vi.fn(),
+			captureTaskCreated: vi.fn(),
+		},
+	},
+}))
+
+// Mock CloudService
+vi.mock("@roo-code/cloud", () => ({
+	CloudService: {
+		hasInstance: vi.fn().mockReturnValue(false),
+		instance: {
+			isAuthenticated: vi.fn().mockReturnValue(false),
+		},
+	},
+	BridgeOrchestrator: {
+		isEnabled: vi.fn().mockReturnValue(false),
+	},
+	getRooCodeApiUrl: vi.fn().mockReturnValue("https://api.roo-code.com"),
+}))
+
+vi.mock("../../../shared/embeddingModels", () => ({
+	EMBEDDING_MODEL_PROFILES: [],
+}))
+
+describe("ClineProvider flicker-free cancel", () => {
+	let provider: ClineProvider
+	let mockContext: any
+	let mockOutputChannel: any
+	let mockTask1: any
+	let mockTask2: any
+
+	const mockApiConfig: ProviderSettings = {
+		apiProvider: "anthropic",
+		apiKey: "test-key",
+	} as ProviderSettings
+
+	beforeEach(() => {
+		vi.clearAllMocks()
+
+		// Setup mock extension context
+		mockContext = {
+			globalState: {
+				get: vi.fn().mockReturnValue(undefined),
+				update: vi.fn().mockResolvedValue(undefined),
+				keys: vi.fn().mockReturnValue([]),
+			},
+			globalStorageUri: { fsPath: "/test/storage" },
+			secrets: {
+				get: vi.fn().mockResolvedValue(undefined),
+				store: vi.fn().mockResolvedValue(undefined),
+				delete: vi.fn().mockResolvedValue(undefined),
+			},
+			workspaceState: {
+				get: vi.fn().mockReturnValue(undefined),
+				update: vi.fn().mockResolvedValue(undefined),
+				keys: vi.fn().mockReturnValue([]),
+			},
+			extensionUri: { fsPath: "/test/extension" },
+		}
+
+		// Setup mock output channel
+		mockOutputChannel = {
+			appendLine: vi.fn(),
+			dispose: vi.fn(),
+		}
+
+		// Setup mock context proxy
+		const mockContextProxy = {
+			getValues: vi.fn().mockReturnValue({}),
+			getValue: vi.fn().mockReturnValue(undefined),
+			setValue: vi.fn().mockResolvedValue(undefined),
+			getProviderSettings: vi.fn().mockReturnValue(mockApiConfig),
+			extensionUri: mockContext.extensionUri,
+			globalStorageUri: mockContext.globalStorageUri,
+		}
+
+		// Create provider instance
+		provider = new ClineProvider(mockContext, mockOutputChannel, "sidebar", mockContextProxy as any)
+
+		// Mock provider methods
+		provider.getState = vi.fn().mockResolvedValue({
+			apiConfiguration: mockApiConfig,
+			mode: "code",
+		})
+
+		provider.postStateToWebview = vi.fn().mockResolvedValue(undefined)
+		// Mock private method using any cast
+		;(provider as any).updateGlobalState = vi.fn().mockResolvedValue(undefined)
+		provider.activateProviderProfile = vi.fn().mockResolvedValue(undefined)
+		provider.performPreparationTasks = vi.fn().mockResolvedValue(undefined)
+		provider.getTaskWithId = vi.fn().mockImplementation((id) =>
+			Promise.resolve({
+				historyItem: {
+					id,
+					number: 1,
+					ts: Date.now(),
+					task: "test task",
+					tokensIn: 100,
+					tokensOut: 200,
+					totalCost: 0.001,
+					workspace: "/test/workspace",
+				},
+			}),
+		)
+
+		// Setup mock tasks
+		mockTask1 = {
+			taskId: "task-1",
+			instanceId: "instance-1",
+			emit: vi.fn(),
+			abortTask: vi.fn().mockResolvedValue(undefined),
+			abandoned: false,
+			dispose: vi.fn(),
+			on: vi.fn(),
+			off: vi.fn(),
+		}
+
+		mockTask2 = {
+			taskId: "task-1", // Same ID for rehydration scenario
+			instanceId: "instance-2", // Different instance
+			emit: vi.fn(),
+			on: vi.fn(),
+			off: vi.fn(),
+		}
+
+		// Mock Task constructor
+		vi.mocked(Task).mockImplementation(() => mockTask2 as any)
+	})
+
+	it("should not remove current task from stack when rehydrating same taskId", async () => {
+		// Setup: Add a task to the stack first
+		;(provider as any).clineStack = [mockTask1]
+
+		// Mock event listeners for cleanup
+		;(provider as any).taskEventListeners = new WeakMap()
+		const mockCleanupFunctions = [vi.fn(), vi.fn()]
+		;(provider as any).taskEventListeners.set(mockTask1, mockCleanupFunctions)
+
+		// Spy on removeClineFromStack to verify it's NOT called
+		const removeClineFromStackSpy = vi.spyOn(provider, "removeClineFromStack")
+
+		// Create history item with same taskId as current task
+		const historyItem: HistoryItem = {
+			id: "task-1", // Same as mockTask1.taskId
+			number: 1,
+			task: "test task",
+			ts: Date.now(),
+			tokensIn: 100,
+			tokensOut: 200,
+			totalCost: 0.001,
+			workspace: "/test/workspace",
+		}
+
+		// Act: Create task with history item (should rehydrate in-place)
+		await provider.createTaskWithHistoryItem(historyItem)
+
+		// Assert: removeClineFromStack should NOT be called
+		expect(removeClineFromStackSpy).not.toHaveBeenCalled()
+
+		// Verify the task was replaced in-place
+		expect((provider as any).clineStack).toHaveLength(1)
+		expect((provider as any).clineStack[0]).toBe(mockTask2)
+
+		// Verify old event listeners were cleaned up
+		expect(mockCleanupFunctions[0]).toHaveBeenCalled()
+		expect(mockCleanupFunctions[1]).toHaveBeenCalled()
+
+		// Verify new task received focus event
+		expect(mockTask2.emit).toHaveBeenCalledWith("taskFocused")
+	})
+
+	it("should remove task from stack when creating different task", async () => {
+		// Setup: Add a task to the stack first
+		;(provider as any).clineStack = [mockTask1]
+
+		// Spy on removeClineFromStack to verify it IS called
+		const removeClineFromStackSpy = vi.spyOn(provider, "removeClineFromStack").mockResolvedValue(undefined)
+
+		// Create history item with different taskId
+		const historyItem: HistoryItem = {
+			id: "task-2", // Different from mockTask1.taskId
+			number: 2,
+			task: "different task",
+			ts: Date.now(),
+			tokensIn: 150,
+			tokensOut: 250,
+			totalCost: 0.002,
+			workspace: "/test/workspace",
+		}
+
+		// Act: Create task with different history item
+		await provider.createTaskWithHistoryItem(historyItem)
+
+		// Assert: removeClineFromStack should be called
+		expect(removeClineFromStackSpy).toHaveBeenCalled()
+	})
+
+	it("should handle empty stack gracefully during rehydration attempt", async () => {
+		// Setup: Empty stack
+		;(provider as any).clineStack = []
+
+		// Spy on removeClineFromStack
+		const removeClineFromStackSpy = vi.spyOn(provider, "removeClineFromStack").mockResolvedValue(undefined)
+
+		// Create history item
+		const historyItem: HistoryItem = {
+			id: "task-1",
+			number: 1,
+			task: "test task",
+			ts: Date.now(),
+			tokensIn: 100,
+			tokensOut: 200,
+			totalCost: 0.001,
+			workspace: "/test/workspace",
+		}
+
+		// Act: Should not error and should call removeClineFromStack
+		await provider.createTaskWithHistoryItem(historyItem)
+
+		// Assert: removeClineFromStack should be called (no current task to rehydrate)
+		expect(removeClineFromStackSpy).toHaveBeenCalled()
+	})
+
+	it("should maintain task stack integrity during flicker-free replacement", async () => {
+		// Setup: Stack with multiple tasks
+		const mockParentTask = {
+			taskId: "parent-task",
+			instanceId: "parent-instance",
+			emit: vi.fn(),
+		}
+
+		;(provider as any).clineStack = [mockParentTask, mockTask1]
+		;(provider as any).taskEventListeners = new WeakMap()
+		;(provider as any).taskEventListeners.set(mockTask1, [vi.fn()])
+
+		// Act: Rehydrate the current (top) task
+		const historyItem: HistoryItem = {
+			id: "task-1",
+			number: 1,
+			task: "test task",
+			ts: Date.now(),
+			tokensIn: 100,
+			tokensOut: 200,
+			totalCost: 0.001,
+			workspace: "/test/workspace",
+		}
+
+		await provider.createTaskWithHistoryItem(historyItem)
+
+		// Assert: Stack should maintain parent task and replace current task
+		expect((provider as any).clineStack).toHaveLength(2)
+		expect((provider as any).clineStack[0]).toBe(mockParentTask)
+		expect((provider as any).clineStack[1]).toBe(mockTask2)
+	})
+})
