test: tighten subtasks cancellation replay

Author: roomote

apps/vscode-e2e/src/fixtures/subtasks.ts

@@ -1,8 +1,11 @@
 import { LLMock } from "@copilotkit/aimock"
 
+import { toolResultContains } from "./tool-result"
+
 export const SUBTASK_PARENT_PROMPT = "SUBTASK_PARENT_CANCELLATION_SMOKE"
 export const SUBTASK_CHILD_PROMPT = "SUBTASK_CHILD_CALCULATOR_SMOKE"
 export const SUBTASK_CHILD_FOLLOWUP_ANSWER = "9"
+const INTERRUPTED_TOOL_RESULT = "Task was interrupted before this tool call could be completed."
 
 export function addSubtaskFixtures(mock: InstanceType<typeof LLMock>) {
 	mock.addFixture({
@@ -44,6 +47,27 @@ export function addSubtaskFixtures(mock: InstanceType<typeof LLMock>) {
 	mock.addFixture({
 		match: {
 			toolCallId: "call_subtasks_child_followup_001",
+			predicate: (req) => toolResultContains(req, "call_subtasks_child_followup_001", [INTERRUPTED_TOOL_RESULT]),
+		},
+		response: {
+			toolCalls: [
+				{
+					name: "ask_followup_question",
+					arguments: JSON.stringify({
+						question: "What is the square root of 81?",
+						follow_up: [{ text: SUBTASK_CHILD_FOLLOWUP_ANSWER }],
+					}),
+					id: "call_subtasks_child_followup_resume_002",
+				},
+			],
+		},
+	})
+
+	mock.addFixture({
+		match: {
+			toolCallId: "call_subtasks_child_followup_001",
+			predicate: (req) =>
+				toolResultContains(req, "call_subtasks_child_followup_001", [SUBTASK_CHILD_FOLLOWUP_ANSWER]),
 		},
 		response: {
 			toolCalls: [
@@ -56,6 +80,23 @@ export function addSubtaskFixtures(mock: InstanceType<typeof LLMock>) {
 		},
 	})
 
+	mock.addFixture({
+		match: {
+			toolCallId: "call_subtasks_child_followup_resume_002",
+			predicate: (req) =>
+				toolResultContains(req, "call_subtasks_child_followup_resume_002", [SUBTASK_CHILD_FOLLOWUP_ANSWER]),
+		},
+		response: {
+			toolCalls: [
+				{
+					name: "attempt_completion",
+					arguments: JSON.stringify({ result: "9" }),
+					id: "call_subtasks_child_completion_resume_003",
+				},
+			],
+		},
+	})
+
 	mock.addFixture({
 		match: {
 			toolCallId: "call_subtasks_parent_new_task_001",

apps/vscode-e2e/src/suite/subtasks.test.ts

@@ -4,7 +4,7 @@ import { RooCodeEventName, type ClineMessage } from "@roo-code/types"
 
 import { setDefaultSuiteTimeout } from "./test-utils"
 import { sleep, waitFor, waitUntilCompleted } from "./utils"
-import { SUBTASK_CHILD_FOLLOWUP_ANSWER, SUBTASK_CHILD_PROMPT, SUBTASK_PARENT_PROMPT } from "../fixtures/subtasks"
+import { SUBTASK_CHILD_FOLLOWUP_ANSWER, SUBTASK_PARENT_PROMPT } from "../fixtures/subtasks"
 
 suite("Roo Code Subtasks", function () {
 	setDefaultSuiteTimeout(this)
@@ -34,6 +34,21 @@ suite("Roo Code Subtasks", function () {
 			}
 		}
 
+		const findCompletionText = (taskId: string) =>
+			messages[taskId]
+				?.filter(
+					(message) =>
+						message.type === "say" && (message.say === "completion_result" || message.say === "text"),
+				)
+				.map((message) => message.text?.trim())
+				.find((text): text is string => !!text)
+
+		const findErrorText = (taskId: string) =>
+			messages[taskId]
+				?.filter((message) => message.type === "say" && message.say === "error")
+				.map((message) => message.text?.trim())
+				.find((text): text is string => !!text)
+
 		api.on(RooCodeEventName.Message, messageHandler)
 
 		try {
@@ -62,6 +77,9 @@ suite("Roo Code Subtasks", function () {
 				"wait for delegated child followup ask",
 				() => asks[spawnedTaskId!]?.some(({ type, ask }) => type === "ask" && ask === "followup") ?? false,
 			)
+			const cancelledChildTaskId = spawnedTaskId!
+			const delegatedFollowupCount =
+				asks[cancelledChildTaskId]?.filter(({ type, ask }) => type === "ask" && ask === "followup").length ?? 0
 
 			await api.cancelCurrentTask()
 
@@ -73,13 +91,41 @@ suite("Roo Code Subtasks", function () {
 				"Parent task should not have resumed after subtask cancellation",
 			)
 
-			const anotherTaskId = await api.startNewTask({ text: SUBTASK_CHILD_PROMPT })
 			await waitForStage(
-				"wait for standalone child followup ask",
-				() => asks[anotherTaskId]?.some(({ type, ask }) => type === "ask" && ask === "followup") ?? false,
+				"wait for cancelled child task to remain active",
+				() => api.getCurrentTaskStack().at(-1) === cancelledChildTaskId,
+			)
+			await waitForStage(
+				"wait for cancelled child resume ask",
+				() =>
+					asks[cancelledChildTaskId]?.some(({ type, ask }) => type === "ask" && ask === "resume_task") ??
+					false,
+			)
+			await api.approveCurrentAsk()
+			await waitForStage(
+				"wait for resumed child followup ask",
+				() =>
+					(asks[cancelledChildTaskId]?.filter(({ type, ask }) => type === "ask" && ask === "followup")
+						.length ?? 0) > delegatedFollowupCount,
 			)
 			await api.sendMessage(SUBTASK_CHILD_FOLLOWUP_ANSWER)
-			await waitUntilCompleted({ api, taskId: anotherTaskId })
+			await waitUntilCompleted({ api, taskId: cancelledChildTaskId })
+
+			assert.strictEqual(
+				findErrorText(cancelledChildTaskId),
+				undefined,
+				"Cancelled child should not emit an error",
+			)
+			assert.strictEqual(
+				findCompletionText(cancelledChildTaskId),
+				"9",
+				"Cancelled child should complete with `9`",
+			)
+			assert.strictEqual(
+				api.getCurrentTaskStack().at(-1),
+				cancelledChildTaskId,
+				"Cancelled child should stay active after resuming from cancellation",
+			)
 
 			await sleep(2_000)
 

src/core/webview/ClineProvider.ts

@@ -2900,8 +2900,8 @@ export class ClineProvider
 		}
 
 		// Preserve parent and root task information for history item.
-		const rootTask = task.rootTask
-		const parentTask = task.parentTask
+		let rootTask = task.rootTask
+		let parentTask = task.parentTask
 
 		// Mark this as a user-initiated cancellation so provider-only rehydration can occur
 		task.abortReason = "user_cancelled"
@@ -2959,6 +2959,34 @@ export class ClineProvider
 			return
 		}
 
+		if (task.parentTaskId) {
+			try {
+				const { historyItem: parentHistory } = await this.getTaskWithId(task.parentTaskId)
+
+				if (parentHistory.status === "delegated" && parentHistory.awaitingChildId === task.taskId) {
+					await this.updateTaskHistory({
+						...parentHistory,
+						status: "active",
+						awaitingChildId: undefined,
+					})
+
+					historyItem = {
+						...historyItem,
+						parentTaskId: undefined,
+						rootTaskId: undefined,
+					}
+					parentTask = undefined
+					rootTask = undefined
+				}
+			} catch (error) {
+				this.log(
+					`[cancelTask] Failed to detach delegated parent for ${task.taskId}: ${
+						error instanceof Error ? error.message : String(error)
+					}`,
+				)
+			}
+		}
+
 		// Clears task again, so we need to abortTask manually above.
 		await this.createTaskWithHistoryItem({ ...historyItem, rootTask, parentTask })
 	}

src/core/webview/__tests__/ClineProvider.flicker-free-cancel.spec.ts

@@ -316,4 +316,82 @@ describe("ClineProvider flicker-free cancel", () => {
 		expect((provider as any).clineStack[0]).toBe(mockParentTask)
 		expect((provider as any).clineStack[1]).toBe(mockTask2)
 	})
+
+	it("detaches a cancelled delegated child before rehydrating it", async () => {
+		const mockRootTask = { taskId: "root-1" }
+		const mockParentTask = { taskId: "parent-1" }
+		const childHistory: HistoryItem = {
+			id: "child-1",
+			number: 2,
+			task: "child task",
+			ts: Date.now(),
+			tokensIn: 10,
+			tokensOut: 20,
+			totalCost: 0.001,
+			workspace: "/test/workspace",
+			parentTaskId: "parent-1",
+			rootTaskId: "root-1",
+		}
+		const parentHistory: HistoryItem = {
+			id: "parent-1",
+			number: 1,
+			task: "parent task",
+			ts: Date.now(),
+			tokensIn: 10,
+			tokensOut: 20,
+			totalCost: 0.001,
+			workspace: "/test/workspace",
+			status: "delegated",
+			awaitingChildId: "child-1",
+			delegatedToId: "child-1",
+		}
+
+		Object.assign(mockTask1, {
+			taskId: "child-1",
+			instanceId: "instance-child",
+			rootTask: mockRootTask,
+			parentTask: mockParentTask,
+			parentTaskId: "parent-1",
+			cancelCurrentRequest: vi.fn(),
+			abortTask: vi.fn(),
+			abandoned: false,
+			isStreaming: false,
+			didFinishAbortingStream: true,
+			isWaitingForFirstChunk: false,
+		})
+		;(provider as any).clineStack = [mockTask1]
+		provider.getTaskWithId = vi.fn().mockImplementation((id) => {
+			if (id === "child-1") {
+				return Promise.resolve({ historyItem: childHistory })
+			}
+			if (id === "parent-1") {
+				return Promise.resolve({ historyItem: parentHistory })
+			}
+			throw new Error(`unexpected task lookup: ${id}`)
+		}) as any
+
+		const updateTaskHistorySpy = vi.spyOn(provider, "updateTaskHistory").mockResolvedValue([])
+		const createTaskWithHistoryItemSpy = vi
+			.spyOn(provider, "createTaskWithHistoryItem")
+			.mockResolvedValue(undefined as any)
+
+		await provider.cancelTask()
+
+		expect(updateTaskHistorySpy).toHaveBeenCalledWith(
+			expect.objectContaining({
+				id: "parent-1",
+				status: "active",
+				awaitingChildId: undefined,
+			}),
+		)
+		expect(createTaskWithHistoryItemSpy).toHaveBeenCalledWith(
+			expect.objectContaining({
+				id: "child-1",
+				parentTaskId: undefined,
+				rootTaskId: undefined,
+				parentTask: undefined,
+				rootTask: undefined,
+			}),
+		)
+	})
 })