Merge pull request #1564 from cannuri/cannuri/fix_subtask_cancel_resume_bug

fix: Preserve parent-child relationship when cancelling subtasks

e2e/src/suite/task.test.ts

@@ -1,4 +1,5 @@
 import * as assert from "assert"
+import * as vscode from "vscode"
 
 suite("Roo Code Task", () => {
 	test("Should handle prompt and response correctly", async function () {
@@ -48,4 +49,113 @@ suite("Roo Code Task", () => {
 			"Did not receive expected response containing 'My name is Roo'",
 		)
 	})
+
+	test("Should handle subtask cancellation and resumption correctly", async function () {
+		this.timeout(60000) // Increase timeout for this test
+		const interval = 1000
+
+		if (!globalThis.extension) {
+			assert.fail("Extension not found")
+		}
+
+		// Ensure the webview is launched
+		await ensureWebviewLaunched(30000, interval)
+
+		// Set up required global state
+		await globalThis.provider.updateGlobalState("mode", "Code")
+		await globalThis.provider.updateGlobalState("alwaysAllowModeSwitch", true)
+		await globalThis.provider.updateGlobalState("alwaysAllowSubtasks", true)
+		await globalThis.provider.updateGlobalState("autoApprovalEnabled", true)
+
+		// 1. Start a parent task that will create a subtask
+		await globalThis.api.startNewTask(
+			"You are the parent task. Create a subtask by using the new_task tool with the message 'You are the subtask'. " +
+				"After creating the subtask, wait for it to complete and then respond with 'Parent task resumed'.",
+		)
+
+		// Wait for the parent task to use the new_task tool
+		await waitForToolUse("new_task", 30000, interval)
+
+		// Wait for the subtask to be created and start responding
+		await waitForMessage("You are the subtask", 10000, interval)
+
+		// 3. Cancel the current task (which should be the subtask)
+		await globalThis.provider.cancelTask()
+
+		// 4. Check if the parent task is still waiting (not resumed)
+		// We need to wait a bit to ensure any task resumption would have happened
+		await new Promise((resolve) => setTimeout(resolve, 5000))
+
+		// The parent task should not have resumed yet, so we shouldn't see "Parent task resumed"
+		assert.ok(
+			!globalThis.provider.messages.some(
+				({ type, text }) => type === "say" && text?.includes("Parent task resumed"),
+			),
+			"Parent task should not have resumed after subtask cancellation",
+		)
+
+		// 5. Start a new task with the same message as the subtask
+		await globalThis.api.startNewTask("You are the subtask")
+
+		// Wait for the subtask to complete
+		await waitForMessage("Task complete", 20000, interval)
+
+		// 6. Verify that the parent task is still not resumed
+		// We need to wait a bit to ensure any task resumption would have happened
+		await new Promise((resolve) => setTimeout(resolve, 5000))
+
+		// The parent task should still not have resumed
+		assert.ok(
+			!globalThis.provider.messages.some(
+				({ type, text }) => type === "say" && text?.includes("Parent task resumed"),
+			),
+			"Parent task should not have resumed after subtask completion",
+		)
+
+		// Clean up - cancel all tasks
+		await globalThis.provider.cancelTask()
+	})
 })
+
+// Helper functions
+async function ensureWebviewLaunched(timeout: number, interval: number): Promise<void> {
+	const startTime = Date.now()
+	while (Date.now() - startTime < timeout) {
+		if (globalThis.provider.viewLaunched) {
+			return
+		}
+		await new Promise((resolve) => setTimeout(resolve, interval))
+	}
+	throw new Error("Webview failed to launch within timeout")
+}
+
+async function waitForToolUse(toolName: string, timeout: number, interval: number): Promise<void> {
+	const startTime = Date.now()
+	while (Date.now() - startTime < timeout) {
+		const messages = globalThis.provider.messages
+		if (
+			messages.some(
+				(message) =>
+					message.type === "say" && message.say === "tool" && message.text && message.text.includes(toolName),
+			)
+		) {
+			return
+		}
+		await new Promise((resolve) => setTimeout(resolve, interval))
+	}
+	throw new Error(`Tool ${toolName} was not used within timeout`)
+}
+
+async function waitForMessage(messageContent: string, timeout: number, interval: number): Promise<void> {
+	const startTime = Date.now()
+	while (Date.now() - startTime < timeout) {
+		const messages = globalThis.provider.messages
+		if (
+			messages.some((message) => message.type === "say" && message.text && message.text.includes(messageContent))
+		) {
+			return
+		}
+		await new Promise((resolve) => setTimeout(resolve, interval))
+	}
+	throw new Error(`Message containing "${messageContent}" not found within timeout`)
+}

src/core/webview/ClineProvider.ts

@@ -2034,8 +2034,16 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 
 	async cancelTask() {
 		if (this.getCurrentCline()) {
-			const { historyItem } = await this.getTaskWithId(this.getCurrentCline()!.taskId)
-			this.getCurrentCline()!.abortTask()
+			const currentCline = this.getCurrentCline()!
+			const { historyItem } = await this.getTaskWithId(currentCline.taskId)
+
+			// Store parent task information if this is a subtask
+			// Check if this is a subtask by seeing if it has a parent task
+			const parentTask = currentCline.getParentTask()
+			const isSubTask = parentTask !== undefined
+			const rootTask = isSubTask ? currentCline.getRootTask() : undefined
+
+			currentCline.abortTask()
 
 			await pWaitFor(
 				() =>
@@ -2062,6 +2070,18 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 
 			// Clears task again, so we need to abortTask manually above.
 			await this.initClineWithHistoryItem(historyItem)
+
+			// Restore parent-child relationship if this was a subtask
+			if (isSubTask && this.getCurrentCline() && parentTask) {
+				this.getCurrentCline()!.setSubTask()
+				this.getCurrentCline()!.setParentTask(parentTask)
+				if (rootTask) {
+					this.getCurrentCline()!.setRootTask(rootTask)
+				}
+				this.log(
+					`[subtasks] Restored parent-child relationship for task: ${this.getCurrentCline()!.getTaskNumber()}`,
+				)
+			}
 		}
 	}