Skip to content

Commit 7cb1196

Browse files
edelaunaedelauna
committed
fix: harden delegated subtask cancellation
1 parent 37c3c9d commit 7cb1196

5 files changed

Lines changed: 79 additions & 70 deletions

File tree

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

Lines changed: 13 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,24 @@ import { toolResultContains } from "./tool-result"
66
const SUBTASK_PARENT_MARKER = "SUBTASK_PARENT_CANCELLATION_SMOKE"
77
const SUBTASK_CHILD_MARKER = "SUBTASK_CHILD_CALCULATOR_SMOKE"
88

9-
const SUBTASK_CHILD_PROMPT = `${SUBTASK_CHILD_MARKER}: Ask the user exactly this follow-up question: What is the square root of 81? After the user answers, complete with only the answer.`
9+
export const SUBTASK_CHILD_PROMPT = `${SUBTASK_CHILD_MARKER}: Ask the user exactly this follow-up question: What is the square root of 81? After the user answers, complete with only the answer.`
1010
export const SUBTASK_PARENT_PROMPT = `${SUBTASK_PARENT_MARKER}: Use the new_task tool exactly once. Create an ask-mode subtask with this exact message: "${SUBTASK_CHILD_PROMPT}" Do not answer directly.`
1111
export const SUBTASK_CHILD_FOLLOWUP_ANSWER = "9"
12-
const INTERRUPTED_TOOL_RESULT = "Task was interrupted before this tool call could be completed."
12+
13+
const requestContains = (req: ChatCompletionRequest, expected: string[]) => {
14+
const rawRequest = JSON.stringify(req)
15+
return expected.every((text) => rawRequest.includes(text))
16+
}
1317

1418
const completionAfterAnswer = (followupId: string, completionId: string) => ({
1519
match: {
16-
toolCallId: followupId,
17-
predicate: (req: ChatCompletionRequest) => toolResultContains(req, followupId, [SUBTASK_CHILD_FOLLOWUP_ANSWER]),
20+
predicate: (req: ChatCompletionRequest) =>
21+
toolResultContains(req, followupId, [SUBTASK_CHILD_FOLLOWUP_ANSWER]) ||
22+
requestContains(req, [followupId, SUBTASK_CHILD_FOLLOWUP_ANSWER]) ||
23+
requestContains(req, [
24+
SUBTASK_CHILD_MARKER,
25+
`<user_message>\\n${SUBTASK_CHILD_FOLLOWUP_ANSWER}\\n</user_message>`,
26+
]),
1827
},
1928
response: {
2029
toolCalls: [
@@ -64,31 +73,8 @@ export function addSubtaskFixtures(mock: InstanceType<typeof LLMock>) {
6473
},
6574
})
6675

67-
mock.addFixture({
68-
match: {
69-
toolCallId: "call_subtasks_child_followup_001",
70-
predicate: (req) => toolResultContains(req, "call_subtasks_child_followup_001", [INTERRUPTED_TOOL_RESULT]),
71-
},
72-
response: {
73-
toolCalls: [
74-
{
75-
name: "ask_followup_question",
76-
arguments: JSON.stringify({
77-
question: "What is the square root of 81?",
78-
follow_up: [{ text: SUBTASK_CHILD_FOLLOWUP_ANSWER }],
79-
}),
80-
id: "call_subtasks_child_followup_resume_002",
81-
},
82-
],
83-
},
84-
})
85-
8676
mock.addFixture(completionAfterAnswer("call_subtasks_child_followup_001", "call_subtasks_child_completion_002"))
8777

88-
mock.addFixture(
89-
completionAfterAnswer("call_subtasks_child_followup_resume_002", "call_subtasks_child_completion_resume_003"),
90-
)
91-
9278
mock.addFixture({
9379
match: {
9480
toolCallId: "call_subtasks_parent_new_task_001",

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

Lines changed: 23 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,13 @@ import * as assert from "assert"
33
import { RooCodeEventName, type ClineMessage } from "@roo-code/types"
44

55
import { setDefaultSuiteTimeout } from "./test-utils"
6-
import { sleep, waitFor, waitUntilCompleted } from "./utils"
7-
import { SUBTASK_CHILD_FOLLOWUP_ANSWER, SUBTASK_PARENT_PROMPT } from "../fixtures/subtasks"
6+
import { waitFor, waitUntilCompleted } from "./utils"
7+
import { SUBTASK_CHILD_FOLLOWUP_ANSWER, SUBTASK_CHILD_PROMPT, SUBTASK_PARENT_PROMPT } from "../fixtures/subtasks"
88

99
suite("Roo Code Subtasks", function () {
1010
setDefaultSuiteTimeout(this)
1111

12-
test("Should handle subtask cancellation and resumption correctly", async () => {
12+
test("Should keep parent paused after subtask cancellation", async () => {
1313
const api = globalThis.api
1414
const asks: Record<string, ClineMessage[]> = {}
1515
const messages: Record<string, ClineMessage[]> = {}
@@ -78,13 +78,9 @@ suite("Roo Code Subtasks", function () {
7878
() => asks[spawnedTaskId!]?.some(({ type, ask }) => type === "ask" && ask === "followup") ?? false,
7979
)
8080
const cancelledChildTaskId = spawnedTaskId!
81-
const delegatedFollowupCount =
82-
asks[cancelledChildTaskId]?.filter(({ type, ask }) => type === "ask" && ask === "followup").length ?? 0
8381

8482
await api.cancelCurrentTask()
8583

86-
await sleep(2_000)
87-
8884
assert.ok(
8985
messages[parentTaskId]?.find(({ type, text }) => type === "say" && text === "Parent task resumed") ===
9086
undefined,
@@ -101,38 +97,40 @@ suite("Roo Code Subtasks", function () {
10197
asks[cancelledChildTaskId]?.some(({ type, ask }) => type === "ask" && ask === "resume_task") ??
10298
false,
10399
)
104-
await api.approveCurrentAsk()
105-
await waitForStage(
106-
"wait for resumed child followup ask",
107-
() =>
108-
(asks[cancelledChildTaskId]?.filter(({ type, ask }) => type === "ask" && ask === "followup")
109-
.length ?? 0) > delegatedFollowupCount,
110-
)
111-
await api.sendMessage(SUBTASK_CHILD_FOLLOWUP_ANSWER)
112-
await waitUntilCompleted({ api, taskId: cancelledChildTaskId })
100+
101+
const standaloneTaskId = await waitUntilCompleted({
102+
api,
103+
start: async () => {
104+
const taskId = await api.startNewTask({ text: SUBTASK_CHILD_PROMPT })
105+
await waitForStage(
106+
"wait for standalone child followup ask",
107+
() => asks[taskId]?.some(({ type, ask }) => type === "ask" && ask === "followup") ?? false,
108+
)
109+
await api.sendMessage(SUBTASK_CHILD_FOLLOWUP_ANSWER)
110+
return taskId
111+
},
112+
})
113113

114114
assert.strictEqual(
115-
findErrorText(cancelledChildTaskId),
115+
findErrorText(standaloneTaskId),
116116
undefined,
117-
"Cancelled child should not emit an error",
117+
"Standalone child task should not emit an error",
118118
)
119119
assert.strictEqual(
120-
findCompletionText(cancelledChildTaskId),
120+
findCompletionText(standaloneTaskId),
121121
"9",
122-
"Cancelled child should complete with `9`",
122+
"Standalone child task should complete with `9`",
123123
)
124124
assert.strictEqual(
125125
api.getCurrentTaskStack().at(-1),
126-
cancelledChildTaskId,
127-
"Cancelled child should stay active after resuming from cancellation",
126+
standaloneTaskId,
127+
"Standalone child task should be the active completed task",
128128
)
129129

130-
await sleep(2_000)
131-
132130
assert.ok(
133131
messages[parentTaskId]?.find(({ type, text }) => type === "say" && text === "Parent task resumed") ===
134132
undefined,
135-
"Parent task should not have resumed after subtask cancellation",
133+
"Parent task should not have resumed after starting another task",
136134
)
137135

138136
await api.clearCurrentTask()

‎src/core/assistant-message/presentAssistantMessage.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ import { sanitizeToolUseId } from "../../utils/tool-id"
6060

6161
export async function presentAssistantMessage(cline: Task) {
6262
if (cline.abort) {
63-
throw new Error(`[Task#presentAssistantMessage] task ${cline.taskId}.${cline.instanceId} aborted`)
63+
return
6464
}
6565

6666
if (cline.presentAssistantMessageLocked) {

‎src/core/webview/ClineProvider.ts‎

Lines changed: 32 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -114,13 +114,11 @@ export type ClineProviderEvents = {
114114
clineCreated: [cline: Task]
115115
}
116116

117-
type DelegationLockHost = {
118-
delegationTransitionLocks?: Map<string, Promise<void>>
119-
}
120-
121-
function runDelegationTransition<T>(host: DelegationLockHost, parentTaskId: string, fn: () => Promise<T>): Promise<T> {
122-
host.delegationTransitionLocks ??= new Map()
123-
const locks = host.delegationTransitionLocks
117+
function runDelegationTransition<T>(
118+
locks: Map<string, Promise<void>>,
119+
parentTaskId: string,
120+
fn: () => Promise<T>,
121+
): Promise<T> {
124122
const previous = locks.get(parentTaskId) ?? Promise.resolve()
125123
const current = previous.then(fn, fn)
126124
const tail = current.then(
@@ -153,7 +151,7 @@ export class ClineProvider
153151
private webviewDisposables: vscode.Disposable[] = []
154152
private view?: vscode.WebviewView | vscode.WebviewPanel
155153
private clineStack: Task[] = []
156-
private delegationTransitionLocks = new Map<string, Promise<void>>()
154+
private delegationTransitionLocks?: Map<string, Promise<void>>
157155
private cancelledDelegationChildIds = new Set<string>()
158156
private codeIndexStatusSubscription?: vscode.Disposable
159157
private codeIndexManager?: CodeIndexManager
@@ -173,6 +171,11 @@ export class ClineProvider
173171
private globalStateWriteThroughTimer: ReturnType<typeof setTimeout> | null = null
174172
private static readonly GLOBAL_STATE_WRITE_THROUGH_DEBOUNCE_MS = 5000 // 5 seconds
175173
private static readonly PENDING_OPERATION_TIMEOUT_MS = 30000 // 30 seconds
174+
175+
private runDelegationTransition<T>(parentTaskId: string, fn: () => Promise<T>): Promise<T> {
176+
this.delegationTransitionLocks ??= new Map()
177+
return runDelegationTransition(this.delegationTransitionLocks, parentTaskId, fn)
178+
}
176179
private readonly pendingEditOperations: PendingEditOperationStore
177180

178181
private cloudOrganizationsCache: CloudOrganizationMembership[] | null = null
@@ -509,7 +512,7 @@ export class ClineProvider
509512
// child and will update the parent to point at the new child.
510513
if (parentTaskId && childTaskId && !options?.skipDelegationRepair) {
511514
try {
512-
await runDelegationTransition(this as unknown as DelegationLockHost, parentTaskId, async () => {
515+
await ClineProvider.prototype.runDelegationTransition.call(this, parentTaskId, async () => {
513516
const { historyItem: parentHistory } = await this.getTaskWithId(parentTaskId)
514517

515518
if (parentHistory?.status === "delegated" && parentHistory?.awaitingChildId === childTaskId) {
@@ -2992,7 +2995,7 @@ export class ClineProvider
29922995

29932996
if (task.parentTaskId) {
29942997
try {
2995-
await runDelegationTransition(this as unknown as DelegationLockHost, task.parentTaskId, async () => {
2998+
await ClineProvider.prototype.runDelegationTransition.call(this, task.parentTaskId, async () => {
29962999
const { historyItem: parentHistory } = await this.getTaskWithId(task.parentTaskId!)
29973000

29983001
if (parentHistory?.status === "delegated" && parentHistory?.awaitingChildId === task.taskId) {
@@ -3011,12 +3014,26 @@ export class ClineProvider
30113014
}
30123015
})
30133016
} catch (error) {
3014-
// Fail closed: if we cannot prove the parent was detached, keep the
3015-
// rehydrated child disconnected from runtime parent links and prevent
3016-
// this in-process child completion from reopening the parent later.
3017+
// Fail closed: if we cannot prove the parent was detached, make the
3018+
// rehydrated child standalone so later completions cannot reopen a
3019+
// stale delegated parent, even after a provider reload.
30173020
parentTask = undefined
30183021
rootTask = undefined
30193022
this.cancelledDelegationChildIds.add(task.taskId)
3023+
historyItem = {
3024+
...historyItem,
3025+
parentTaskId: undefined,
3026+
rootTaskId: undefined,
3027+
}
3028+
try {
3029+
await this.updateTaskHistory(historyItem)
3030+
} catch (historyError) {
3031+
this.log(
3032+
`[cancelTask] Failed to persist standalone child state for ${task.taskId}: ${
3033+
historyError instanceof Error ? historyError.message : String(historyError)
3034+
}`,
3035+
)
3036+
}
30203037
this.log(
30213038
`[cancelTask] Failed to detach delegated parent for ${task.taskId}: ${
30223039
error instanceof Error ? error.message : String(error)
@@ -3379,7 +3396,7 @@ export class ClineProvider
33793396
completionResultSummary: string
33803397
}): Promise<boolean> {
33813398
const { parentTaskId, childTaskId, completionResultSummary } = params
3382-
return await runDelegationTransition(this as unknown as DelegationLockHost, parentTaskId, async () => {
3399+
return await (ClineProvider.prototype.runDelegationTransition.call(this, parentTaskId, async () => {
33833400
const globalStoragePath = this.contextProxy.globalStorageUri.fsPath
33843401

33853402
// 1) Load parent from history and current persisted messages
@@ -3588,7 +3605,7 @@ export class ClineProvider
35883605

35893606
;(this.cancelledDelegationChildIds as Set<string> | undefined)?.delete(childTaskId)
35903607
return true
3591-
})
3608+
}) as Promise<boolean>)
35923609
}
35933610

35943611
/**

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -435,6 +435,7 @@ describe("ClineProvider flicker-free cancel", () => {
435435
throw new Error(`unexpected task lookup: ${id}`)
436436
}) as any
437437

438+
const updateTaskHistorySpy = vi.spyOn(provider, "updateTaskHistory").mockResolvedValue([])
438439
const createTaskWithHistoryItemSpy = vi
439440
.spyOn(provider, "createTaskWithHistoryItem")
440441
.mockResolvedValue(undefined as any)
@@ -444,11 +445,18 @@ describe("ClineProvider flicker-free cancel", () => {
444445
expect(mockOutputChannel.appendLine).toHaveBeenCalledWith(
445446
expect.stringContaining("[cancelTask] Failed to detach delegated parent for child-1: parent lookup failed"),
446447
)
448+
expect(updateTaskHistorySpy).toHaveBeenCalledWith(
449+
expect.objectContaining({
450+
id: "child-1",
451+
parentTaskId: undefined,
452+
rootTaskId: undefined,
453+
}),
454+
)
447455
expect(createTaskWithHistoryItemSpy).toHaveBeenCalledWith(
448456
expect.objectContaining({
449457
id: "child-1",
450-
parentTaskId: "parent-1",
451-
rootTaskId: "root-1",
458+
parentTaskId: undefined,
459+
rootTaskId: undefined,
452460
parentTask: undefined,
453461
rootTask: undefined,
454462
}),

0 commit comments

Comments
 (0)