Skip to content

Commit 43be7e1

Browse files
PierrunoYTPierrunoYT
committed
fix(delegation): serialize rehydration install and re-delegation commit with task ownership reservation
1 parent 7af3026 commit 43be7e1

3 files changed

Lines changed: 141 additions & 20 deletions

File tree

‎src/__tests__/ClineProvider.delegation.spec.ts‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,7 @@ describe("ClineProvider.delegateParentAndOpenChild()", () => {
177177
deleteTaskWithId,
178178
createTaskWithHistoryItem,
179179
log: vi.fn(),
180+
isTaskRunningInAnyProvider: vi.fn().mockReturnValue(false),
180181
isViewLaunched: false,
181182
} as unknown as ClineProvider
182183
Reflect.set(provider, "taskHistoryStore", {
@@ -645,6 +646,7 @@ describe("ClineProvider.delegateParentAndOpenChild()", () => {
645646
createTask: vi.fn().mockResolvedValue({ taskId: "child-2", start: vi.fn(), run: () => Promise.resolve() }),
646647
handleModeSwitch: vi.fn().mockResolvedValue(undefined),
647648
log: vi.fn(),
649+
isTaskRunningInAnyProvider: vi.fn().mockReturnValue(false),
648650
isViewLaunched: false,
649651
recentTasksCache: undefined,
650652
taskHistoryStore,
@@ -908,6 +910,67 @@ describe("ClineProvider.delegateParentAndOpenChild()", () => {
908910
expect(taskRegistry.push).toHaveBeenCalledOnce()
909911
})
910912

913+
it("rejects re-delegation when the interrupted awaited child becomes live before the commit", async () => {
914+
const oldChildId = "old-child"
915+
const parent: HistoryItem = {
916+
...parentHistoryItem,
917+
status: "delegated",
918+
awaitingChildId: oldChildId,
919+
delegatedToId: oldChildId,
920+
childIds: [oldChildId],
921+
}
922+
const records = new Map<string, HistoryItem>([
923+
["parent-1", parent],
924+
[oldChildId, { ...parentHistoryItem, id: oldChildId, status: "interrupted", parentTaskId: "parent-1" }],
925+
])
926+
let oldChildIsLive = false
927+
const taskHistoryStore = makeStoreStub({
928+
get: vi.fn((id: string) => records.get(id)),
929+
atomicReadAndUpdate: vi.fn(async (taskId: string, updater: (item: HistoryItem) => HistoryItem) => {
930+
records.set(taskId, updater(records.get(taskId)!))
931+
return []
932+
}),
933+
})
934+
const parentTask = makeParentTask()
935+
const provider = {
936+
_disposed: false,
937+
taskScheduler: new TaskScheduler(),
938+
taskRegistry: { hasRunning: vi.fn((id: string) => id === oldChildId && oldChildIsLive) },
939+
emit: vi.fn(),
940+
getCurrentTask: vi.fn(() => parentTask),
941+
removeClineFromStack: vi.fn().mockResolvedValue(undefined),
942+
// The old child is resumed elsewhere after the pre-check but before the parent commit.
943+
createTask: vi.fn(async () => {
944+
oldChildIsLive = true
945+
return { taskId: "child-2", run: vi.fn().mockResolvedValue(undefined) }
946+
}),
947+
deleteTaskWithId: vi.fn().mockResolvedValue(undefined),
948+
getTaskWithId: vi.fn(async () => ({ historyItem: records.get("parent-1") })),
949+
createTaskWithHistoryItem: vi.fn().mockResolvedValue(undefined),
950+
log: vi.fn(),
951+
isViewLaunched: false,
952+
recentTasksCache: undefined,
953+
taskHistoryStore,
954+
} as unknown as ClineProvider
955+
Object.assign(provider, {
956+
isTaskRunningInAnyProvider: ClineProvider.prototype["isTaskRunningInAnyProvider"],
957+
})
958+
959+
await expect(
960+
ClineProvider.prototype.delegateParentAndOpenChild.call(provider, {
961+
parentTaskId: "parent-1",
962+
message: "Continue",
963+
initialTodos: [],
964+
mode: "code",
965+
}),
966+
).rejects.toThrow(`awaited child ${oldChildId} became live before delegation commit`)
967+
expect(records.get("parent-1")).toMatchObject({
968+
status: "delegated",
969+
awaitingChildId: oldChildId,
970+
delegatedToId: oldChildId,
971+
})
972+
})
973+
911974
it("does not recover when a descendant is live in another provider", async () => {
912975
const oldChildId = "old-child"
913976
const grandchildId = "grandchild"

‎src/core/webview/ClineProvider.ts‎

Lines changed: 43 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1661,8 +1661,18 @@ export class ClineProvider
16611661

16621662
removeTaskEventListeners(this.taskEventListeners, oldTask)
16631663

1664-
// Replace in-place: preserves stack index and current pointer
1665-
this.taskRegistry.replace(oldTask.taskId, task)
1664+
// Replace in-place (preserves stack index and current pointer). Serialized with
1665+
// dead-delegation recovery so its persisted write cannot interleave with this install.
1666+
const installed = await withTaskOwnershipReservation(async () => {
1667+
if (this._disposed || this.taskRegistry.current !== oldTask) return false
1668+
this.taskRegistry.replace(oldTask.taskId, task)
1669+
return true
1670+
})
1671+
if (!installed) {
1672+
removeTaskEventListeners(this.taskEventListeners, task)
1673+
await task.abortTask(true).catch(() => {})
1674+
return task
1675+
}
16661676
}
16671677

16681678
task.emit(RooCodeEventName.TaskFocused)
@@ -4991,24 +5001,37 @@ export class ClineProvider
49915001
if (this._disposed) {
49925002
throw new Error("[delegateParentAndOpenChild] Provider was disposed before delegation commit")
49935003
}
4994-
await this.taskHistoryStore.atomicReadAndUpdate(parentTaskId, (historyItem) => {
4995-
if (pendingActionId && historyItem.pendingAction?.actionId !== pendingActionId) {
4996-
throw new Error(
4997-
`[delegateParentAndOpenChild] Pending action mismatch for parent ${parentTaskId}: expected ${pendingActionId}, found ${historyItem.pendingAction?.actionId}`,
4998-
)
4999-
}
5000-
const awaitedChildStatus = historyItem.awaitingChildId
5001-
? this.taskHistoryStore.get(historyItem.awaitingChildId)?.status
5002-
: undefined
5003-
const delegated = delegateTaskToChild(historyItem, child.taskId, awaitedChildStatus)
5004-
pendingActionBeforeCommit = historyItem.pendingAction
5005-
parentBeforeCommit = historyItem
5006-
return {
5007-
...delegated,
5008-
pendingAction:
5009-
delegated.pendingAction?.actionId === pendingActionId ? undefined : delegated.pendingAction,
5010-
}
5011-
})
5004+
// Hold the ownership reservation across the commit so the awaited child cannot register
5005+
// between the liveness check and the write; only the commit is reserved (non-reentrant).
5006+
await withTaskOwnershipReservation(() =>
5007+
this.taskHistoryStore.atomicReadAndUpdate(parentTaskId, (historyItem) => {
5008+
if (pendingActionId && historyItem.pendingAction?.actionId !== pendingActionId) {
5009+
throw new Error(
5010+
`[delegateParentAndOpenChild] Pending action mismatch for parent ${parentTaskId}: expected ${pendingActionId}, found ${historyItem.pendingAction?.actionId}`,
5011+
)
5012+
}
5013+
if (
5014+
historyItem.status === "delegated" &&
5015+
historyItem.awaitingChildId &&
5016+
this.isTaskRunningInAnyProvider(historyItem.awaitingChildId)
5017+
) {
5018+
throw new LifecycleTransitionError(
5019+
`Cannot re-delegate task ${parentTaskId}: awaited child ${historyItem.awaitingChildId} became live before delegation commit`,
5020+
)
5021+
}
5022+
const awaitedChildStatus = historyItem.awaitingChildId
5023+
? this.taskHistoryStore.get(historyItem.awaitingChildId)?.status
5024+
: undefined
5025+
const delegated = delegateTaskToChild(historyItem, child.taskId, awaitedChildStatus)
5026+
pendingActionBeforeCommit = historyItem.pendingAction
5027+
parentBeforeCommit = historyItem
5028+
return {
5029+
...delegated,
5030+
pendingAction:
5031+
delegated.pendingAction?.actionId === pendingActionId ? undefined : delegated.pendingAction,
5032+
}
5033+
}),
5034+
)
50125035
parentCommitted = true
50135036
if (this._disposed) {
50145037
throw new Error("[delegateParentAndOpenChild] Provider was disposed before child scheduling")

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { ClineProvider } from "../ClineProvider"
77
import { Task } from "../../task/Task"
88
import { TaskRegistry } from "../../task/TaskRegistry"
99
import { ContextProxy } from "../../config/ContextProxy"
10+
import { withTaskOwnershipReservation } from "../../task-persistence"
1011
import type { ProviderSettings, HistoryItem } from "@roo-code/types"
1112
import { providerIdentifiers } from "@roo-code/types/provider-identifiers"
1213

@@ -463,6 +464,40 @@ describe("ClineProvider flicker-free cancel", () => {
463464
expect(mockTask2.emit).toHaveBeenCalledWith("taskFocused")
464465
})
465466

467+
it("serializes the in-place rehydration install with dead-delegation recovery", async () => {
468+
seedRegistry(provider, mockTask1)
469+
let release!: () => void
470+
const reservation = withTaskOwnershipReservation(
471+
() =>
472+
new Promise<void>((resolve) => {
473+
release = resolve
474+
}),
475+
)
476+
await vi.waitFor(() => expect(release).toBeDefined())
477+
478+
const rehydration = provider.createTaskWithHistoryItem({
479+
id: "task-1",
480+
number: 1,
481+
task: "test task",
482+
ts: Date.now(),
483+
tokensIn: 0,
484+
tokensOut: 0,
485+
totalCost: 0,
486+
workspace: "/test/workspace",
487+
})
488+
try {
489+
await vi.waitFor(() => expect(mockTask1.abortTask).toHaveBeenCalled())
490+
await new Promise((resolve) => setTimeout(resolve, 0))
491+
// The recovery reservation is held: the aborted task must not be replaced yet.
492+
expect(provider["taskRegistry"].current).toBe(mockTask1)
493+
} finally {
494+
release()
495+
}
496+
await reservation
497+
await rehydration
498+
expect(provider["taskRegistry"].current).toBe(mockTask2)
499+
})
500+
466501
it("should remove task from stack when creating different task", async () => {
467502
// Setup: Add a task to the registry first
468503
seedRegistry(provider, mockTask1)

0 commit comments

Comments
 (0)