Skip to content

Commit 16fced1

Browse files
edelaunaedelauna
authored
refactor(task): replace dispatch lock booleans with DispatchState (#2198)
1 parent 4d14dbe commit 16fced1

14 files changed

Lines changed: 213 additions & 73 deletions

‎docs/architecture/task-lifecycle-model.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ The model maps its state to production fields as follows:
102102
| -------------------------- | --------------------------------------------------------------------------------------------------------------------------------------- |
103103
| `phase`, `index` | `Task.assistantMessageContent` partial flags, matching `tool_result` blocks in `userMessageContent`, and `currentStreamingContentIndex` |
104104
| `stream` | `didCompleteReadingStream` |
105-
| `lock`, `owner`, `pending` | `presentAssistantMessageLocked` and `presentAssistantMessageHasPendingUpdates`, owned by the `presentAssistantMessage` wrapper |
105+
| `lock`, `owner`, `pending` | `dispatchState` (not `IDLE`) and `pendingWork` (not `null`), owned by the `presentAssistantMessage` wrapper |
106106
| `latch` | `userMessageContentReady` |
107107
| `epoch` | the per-request reset in the request loop (`Task.ts`) |
108108
| `poll` and `post-wait` | the `pWaitFor` condition and the abort check that follows it |

‎scripts/check-tool-turn-continuation.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,9 @@ interface ModelState {
2020
phase: Phase[] // Task.assistantMessageContent partial flag plus tool_result presence
2121
stream: "open" | "complete" // Task.didCompleteReadingStream
2222
index: number // Task.currentStreamingContentIndex
23-
lock: boolean // Task.presentAssistantMessageLocked
23+
lock: boolean // Task.dispatchState !== DispatchState.IDLE
2424
owner: number
25-
pending: boolean // Task.presentAssistantMessageHasPendingUpdates
25+
pending: boolean // Task.pendingWork !== null
2626
latch: boolean // Task.userMessageContentReady
2727
abort: boolean
2828
threw: boolean

‎src/core/assistant-message/__tests__/presentAssistantMessage-auto-deny.spec.ts‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// npx vitest run core/assistant-message/__tests__/presentAssistantMessage-auto-deny.spec.ts
22

3+
import { DispatchState, type PendingWork } from "../../task/dispatchState"
34
import type { Anthropic } from "@anthropic-ai/sdk"
45
import { describe, it, expect, beforeEach, vi } from "vitest"
56
import { presentAssistantMessage } from "../presentAssistantMessage"
@@ -69,8 +70,8 @@ interface MockTask {
6970
taskId: string
7071
instanceId: string
7172
abort: boolean
72-
presentAssistantMessageLocked: boolean
73-
presentAssistantMessageHasPendingUpdates: boolean
73+
dispatchState: DispatchState
74+
pendingWork: PendingWork | null
7475
currentStreamingContentIndex: number
7576
assistantMessageContent: unknown[]
7677
userMessageContent: Anthropic.ToolResultBlockParam[]
@@ -102,8 +103,8 @@ function buildMockTask(): MockTask {
102103
taskId: "test-task-id",
103104
instanceId: "test-instance",
104105
abort: false,
105-
presentAssistantMessageLocked: false,
106-
presentAssistantMessageHasPendingUpdates: false,
106+
dispatchState: DispatchState.IDLE,
107+
pendingWork: null,
107108
currentStreamingContentIndex: 0,
108109
assistantMessageContent: [],
109110
userMessageContent: [],

‎src/core/assistant-message/__tests__/presentAssistantMessage-custom-tool.spec.ts‎

Lines changed: 33 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-custom-tool.spec.ts
22

3+
import { DispatchState } from "../../task/dispatchState"
34
import { describe, it, expect, beforeEach, vi } from "vitest"
45
import { presentAssistantMessage } from "../presentAssistantMessage"
56
import { validateToolUse } from "../../tools/validateToolUse"
@@ -67,8 +68,8 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
6768
taskId: "test-task-id",
6869
instanceId: "test-instance",
6970
abort: false,
70-
presentAssistantMessageLocked: false,
71-
presentAssistantMessageHasPendingUpdates: false,
71+
dispatchState: DispatchState.IDLE,
72+
pendingWork: null,
7273
currentStreamingContentIndex: 0,
7374
assistantMessageContent: [],
7475
userMessageContent: [],
@@ -224,17 +225,17 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
224225
})
225226
})
226227

227-
describe("Presentation lock ownership", () => {
228+
describe("Dispatch state ownership", () => {
228229
it("drains an update queued during the helper-to-wrapper handoff", async () => {
229230
mockTask.assistantMessageContent = [{ type: "text", content: "initial", partial: true }]
230231
let queuedCall: Promise<void> | undefined
231-
let lockedDuringHandoff = false
232+
let busyDuringHandoff = false
232233
mockTask.say.mockImplementationOnce(async () => {
233234
// The first microtask precedes the helper's await continuation.
234235
// The second follows that continuation but precedes the wrapper's.
235236
queueMicrotask(() => {
236237
queueMicrotask(() => {
237-
lockedDuringHandoff = mockTask.presentAssistantMessageLocked
238+
busyDuringHandoff = mockTask.dispatchState !== DispatchState.IDLE
238239
mockTask.assistantMessageContent[0] = { type: "text", content: "final", partial: false }
239240
mockTask.didCompleteReadingStream = true
240241
queuedCall = presentAssistantMessage(mockTask, baseSnapshot)
@@ -245,15 +246,15 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
245246
await presentAssistantMessage(mockTask, baseSnapshot)
246247
await queuedCall
247248

248-
expect(lockedDuringHandoff).toBe(true)
249+
expect(busyDuringHandoff).toBe(true)
249250
expect(mockTask.say.mock.calls).toEqual([
250251
["text", "initial", undefined, true],
251252
["text", "final", undefined, false],
252253
])
253254
expect(mockTask.currentStreamingContentIndex).toBe(1)
254255
expect(mockTask.userMessageContentReady).toBe(true)
255-
expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(false)
256-
expect(mockTask.presentAssistantMessageLocked).toBe(false)
256+
expect(mockTask.pendingWork).toBeNull()
257+
expect(mockTask.dispatchState).toBe(DispatchState.IDLE)
257258
})
258259

259260
it("stops draining an update queued during the handoff once the task aborts", async () => {
@@ -279,7 +280,7 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
279280
mockTask.assistantMessageContent[0] = { type: "text", content: "final", partial: false }
280281
mockTask.didCompleteReadingStream = true
281282
queuedCall = presentAssistantMessage(mockTask, baseSnapshot)
282-
pendingDuringHandoff = mockTask.presentAssistantMessageHasPendingUpdates
283+
pendingDuringHandoff = mockTask.pendingWork !== null
283284
aborted = true
284285
})
285286
})
@@ -291,30 +292,30 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
291292
expect(pendingDuringHandoff).toBe(true)
292293
expect(mockTask.say).toHaveBeenCalledExactlyOnceWith("text", "initial", undefined, true)
293294
expect(mockTask.currentStreamingContentIndex).toBe(0)
294-
expect(mockTask.presentAssistantMessageLocked).toBe(false)
295+
expect(mockTask.dispatchState).toBe(DispatchState.IDLE)
295296
})
296297

297-
it("holds one lock across consecutive blocks and queues overlapping calls", async () => {
298+
it("holds SERIAL across consecutive blocks and queues overlapping calls", async () => {
298299
mockTask.assistantMessageContent = [
299300
{ type: "text", content: "first", partial: false },
300301
{ type: "text", content: "second", partial: false },
301302
]
302303
mockTask.didCompleteReadingStream = true
303-
const lockChanges: boolean[] = []
304-
let locked = false
305-
Object.defineProperty(mockTask, "presentAssistantMessageLocked", {
304+
const stateChanges: DispatchState[] = []
305+
let state = DispatchState.IDLE
306+
Object.defineProperty(mockTask, "dispatchState", {
306307
configurable: true,
307-
get: () => locked,
308-
set: (value: boolean) => {
309-
locked = value
310-
lockChanges.push(value)
308+
get: () => state,
309+
set: (value: DispatchState) => {
310+
state = value
311+
stateChanges.push(value)
311312
},
312313
})
313314
mockTask.say.mockImplementation(async (_type: string, content: string) => {
314-
expect(mockTask.presentAssistantMessageLocked).toBe(true)
315+
expect(mockTask.dispatchState).toBe(DispatchState.SERIAL)
315316
if (content === "second") {
316317
await presentAssistantMessage(mockTask, baseSnapshot)
317-
expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(true)
318+
expect(mockTask.pendingWork).toEqual({ mode: "serial" })
318319
}
319320
})
320321

@@ -326,8 +327,8 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
326327
])
327328
expect(mockTask.currentStreamingContentIndex).toBe(2)
328329
expect(mockTask.userMessageContentReady).toBe(true)
329-
expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(false)
330-
expect(lockChanges).toEqual([true, false])
330+
expect(mockTask.pendingWork).toBeNull()
331+
expect(stateChanges).toEqual([DispatchState.SERIAL, DispatchState.IDLE])
331332
})
332333

333334
it("consumes a pending update for a partial block without repeating it indefinitely", async () => {
@@ -336,11 +337,11 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
336337
.mockImplementationOnce(async () => {
337338
mockTask.assistantMessageContent[0] = { type: "text", content: "updated", partial: true }
338339
await presentAssistantMessage(mockTask, baseSnapshot)
339-
expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(true)
340+
expect(mockTask.pendingWork).toEqual({ mode: "serial" })
340341
})
341342
.mockImplementationOnce(async () => {
342-
expect(mockTask.presentAssistantMessageLocked).toBe(true)
343-
expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(false)
343+
expect(mockTask.dispatchState).toBe(DispatchState.SERIAL)
344+
expect(mockTask.pendingWork).toBeNull()
344345
})
345346

346347
await presentAssistantMessage(mockTask, baseSnapshot)
@@ -350,8 +351,8 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
350351
["text", "updated", undefined, true],
351352
])
352353
expect(mockTask.currentStreamingContentIndex).toBe(0)
353-
expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(false)
354-
expect(mockTask.presentAssistantMessageLocked).toBe(false)
354+
expect(mockTask.pendingWork).toBeNull()
355+
expect(mockTask.dispatchState).toBe(DispatchState.IDLE)
355356
})
356357

357358
it("stops internal continuation when the task aborts between blocks", async () => {
@@ -367,10 +368,10 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
367368

368369
expect(mockTask.say).toHaveBeenCalledExactlyOnceWith("text", "first", undefined, false)
369370
expect(mockTask.currentStreamingContentIndex).toBe(1)
370-
expect(mockTask.presentAssistantMessageLocked).toBe(false)
371+
expect(mockTask.dispatchState).toBe(DispatchState.IDLE)
371372
})
372373

373-
it("releases the presentation lock when a later block throws", async () => {
374+
it("returns to IDLE when a later block throws", async () => {
374375
mockTask.assistantMessageContent = [
375376
{ type: "text", content: "first", partial: false },
376377
{ type: "text", content: "second", partial: false },
@@ -381,12 +382,12 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
381382

382383
expect(mockTask.say).toHaveBeenCalledTimes(2)
383384
expect(mockTask.currentStreamingContentIndex).toBe(1)
384-
expect(mockTask.presentAssistantMessageLocked).toBe(false)
385+
expect(mockTask.dispatchState).toBe(DispatchState.IDLE)
385386
})
386387
})
387388

388389
describe("Custom tool error recording", () => {
389-
it("releases the presentation lock when dispatch throws", async () => {
390+
it("returns to IDLE when dispatch throws", async () => {
390391
mockTask.assistantMessageContent = [
391392
{
392393
type: "tool_use",
@@ -404,7 +405,7 @@ describe("presentAssistantMessage - Custom Tool Recording", () => {
404405

405406
try {
406407
await expect(presentAssistantMessage(mockTask, baseSnapshot)).rejects.toThrow("dispatch failed")
407-
expect(mockTask.presentAssistantMessageLocked).toBe(false)
408+
expect(mockTask.dispatchState).toBe(DispatchState.IDLE)
408409
} finally {
409410
handleSpy.mockRestore()
410411
}

‎src/core/assistant-message/__tests__/presentAssistantMessage-images.spec.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-images.spec.ts
22

3+
import { DispatchState } from "../../task/dispatchState"
34
import { describe, it, expect, beforeEach, vi } from "vitest"
45
import { presentAssistantMessage } from "../presentAssistantMessage"
56
import { validateToolUse } from "../../tools/validateToolUse"
@@ -40,8 +41,8 @@ describe("presentAssistantMessage - Image Handling in Native Tool Calling", () =
4041
taskId: "test-task-id",
4142
instanceId: "test-instance",
4243
abort: false,
43-
presentAssistantMessageLocked: false,
44-
presentAssistantMessageHasPendingUpdates: false,
44+
dispatchState: DispatchState.IDLE,
45+
pendingWork: null,
4546
currentStreamingContentIndex: 0,
4647
assistantMessageContent: [],
4748
userMessageContent: [],

‎src/core/assistant-message/__tests__/presentAssistantMessage-tool-repetition.spec.ts‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-tool-repetition.spec.ts
22

3+
import { DispatchState } from "../../task/dispatchState"
34
import { providerIdentifiers } from "@roo-code/types"
45

56
import type { ToolUse } from "../../../shared/tools"
@@ -153,8 +154,8 @@ describe("presentAssistantMessage - Tool Repetition Detection", () => {
153154
taskId: "test-task-id",
154155
instanceId: "test-instance",
155156
abort: false,
156-
presentAssistantMessageLocked: false,
157-
presentAssistantMessageHasPendingUpdates: false,
157+
dispatchState: DispatchState.IDLE,
158+
pendingWork: null,
158159
currentStreamingContentIndex: 0,
159160
assistantMessageContent: [],
160161
userMessageContent: [],
@@ -217,8 +218,8 @@ describe("presentAssistantMessage - Tool Repetition Detection", () => {
217218
*/
218219
function startNextTurn(blocks: ToolUse[]) {
219220
mockTask.currentStreamingContentIndex = 0
220-
mockTask.presentAssistantMessageLocked = false
221-
mockTask.presentAssistantMessageHasPendingUpdates = false
221+
mockTask.dispatchState = DispatchState.IDLE
222+
mockTask.pendingWork = null
222223
mockTask.didAlreadyUseTool = false
223224
mockTask.didRejectTool = false
224225
mockTask.userMessageContent = []

‎src/core/assistant-message/__tests__/presentAssistantMessage-tool-usage-attribution.spec.ts‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-tool-usage-attribution.spec.ts
22

3+
import { DispatchState, type PendingWork } from "../../task/dispatchState"
34
import type { Anthropic } from "@anthropic-ai/sdk"
45
import { TelemetryEventName } from "@roo-code/types"
56
import { describe, it, expect, beforeEach, vi, type Mock } from "vitest"
@@ -60,8 +61,8 @@ interface MockTask {
6061
taskId: string
6162
instanceId: string
6263
abort: boolean
63-
presentAssistantMessageLocked: boolean
64-
presentAssistantMessageHasPendingUpdates: boolean
64+
dispatchState: DispatchState
65+
pendingWork: PendingWork | null
6566
currentStreamingContentIndex: number
6667
assistantMessageContent: unknown[]
6768
userMessageContent: Anthropic.ToolResultBlockParam[]
@@ -100,8 +101,8 @@ describe("presentAssistantMessage - tool usage attribution", () => {
100101
taskId: "test-task-id",
101102
instanceId: "test-instance",
102103
abort: false,
103-
presentAssistantMessageLocked: false,
104-
presentAssistantMessageHasPendingUpdates: false,
104+
dispatchState: DispatchState.IDLE,
105+
pendingWork: null,
105106
currentStreamingContentIndex: 0,
106107
assistantMessageContent: [],
107108
userMessageContent: [],

‎src/core/assistant-message/__tests__/presentAssistantMessage-unknown-tool.spec.ts‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-unknown-tool.spec.ts
22

3+
import { DispatchState } from "../../task/dispatchState"
34
import { describe, it, expect, beforeEach, vi } from "vitest"
45
import { presentAssistantMessage } from "../presentAssistantMessage"
56
import { isValidToolName, validateToolUse } from "../../tools/validateToolUse"
@@ -42,8 +43,8 @@ describe("presentAssistantMessage - Unknown Tool Handling", () => {
4243
taskId: "test-task-id",
4344
instanceId: "test-instance",
4445
abort: false,
45-
presentAssistantMessageLocked: false,
46-
presentAssistantMessageHasPendingUpdates: false,
46+
dispatchState: DispatchState.IDLE,
47+
pendingWork: null,
4748
currentStreamingContentIndex: 0,
4849
assistantMessageContent: [],
4950
userMessageContent: [],
@@ -243,7 +244,7 @@ describe("presentAssistantMessage - Unknown Tool Handling", () => {
243244
expect(mockTask.say).toHaveBeenCalledTimes(2)
244245
expect(mockTask.say).toHaveBeenNthCalledWith(1, "text", "streaming", undefined, true)
245246
expect(mockTask.say).toHaveBeenNthCalledWith(2, "text", "updated", undefined, true)
246-
expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(false)
247+
expect(mockTask.pendingWork).toBeNull()
247248
expect(mockTask.currentStreamingContentIndex).toBe(0)
248249
})
249250

0 commit comments

Comments
 (0)