Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 46 additions & 5 deletions src/core/webview/ClineProvider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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}`
Expand Down
321 changes: 321 additions & 0 deletions src/core/webview/__tests__/ClineProvider.flicker-free-cancel.spec.ts
Original file line number Diff line number Diff line change
@@ -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)
})
})