From 529861cf19ab063ee45ebc767cd6d6745f86676d Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Sun, 25 Oct 2026 23:38:13 -0400 Subject: [PATCH] refactor(task): replace dispatch lock booleans with DispatchState --- docs/architecture/task-lifecycle-model.md | 2 +- scripts/check-tool-turn-continuation.ts | 4 +- .../presentAssistantMessage-auto-deny.spec.ts | 9 +- ...resentAssistantMessage-custom-tool.spec.ts | 65 +++++------ .../presentAssistantMessage-images.spec.ts | 5 +- ...ntAssistantMessage-tool-repetition.spec.ts | 9 +- ...tantMessage-tool-usage-attribution.spec.ts | 9 +- ...esentAssistantMessage-unknown-tool.spec.ts | 7 +- .../__tests__/presentAssistantMessage.spec.ts | 108 ++++++++++++++++++ .../presentAssistantMessage.ts | 13 ++- src/core/task/Task.ts | 11 +- src/core/task/__tests__/Task.spec.ts | 17 +-- .../__tests__/ToolExecutionContext.spec.ts | 7 +- src/core/task/dispatchState.ts | 20 ++++ 14 files changed, 213 insertions(+), 73 deletions(-) create mode 100644 src/core/assistant-message/__tests__/presentAssistantMessage.spec.ts create mode 100644 src/core/task/dispatchState.ts diff --git a/docs/architecture/task-lifecycle-model.md b/docs/architecture/task-lifecycle-model.md index eaeb61d987..b7b39e600d 100644 --- a/docs/architecture/task-lifecycle-model.md +++ b/docs/architecture/task-lifecycle-model.md @@ -102,7 +102,7 @@ The model maps its state to production fields as follows: | -------------------------- | --------------------------------------------------------------------------------------------------------------------------------------- | | `phase`, `index` | `Task.assistantMessageContent` partial flags, matching `tool_result` blocks in `userMessageContent`, and `currentStreamingContentIndex` | | `stream` | `didCompleteReadingStream` | -| `lock`, `owner`, `pending` | `presentAssistantMessageLocked` and `presentAssistantMessageHasPendingUpdates`, owned by the `presentAssistantMessage` wrapper | +| `lock`, `owner`, `pending` | `dispatchState` (not `IDLE`) and `pendingWork` (not `null`), owned by the `presentAssistantMessage` wrapper | | `latch` | `userMessageContentReady` | | `epoch` | the per-request reset in the request loop (`Task.ts`) | | `poll` and `post-wait` | the `pWaitFor` condition and the abort check that follows it | diff --git a/scripts/check-tool-turn-continuation.ts b/scripts/check-tool-turn-continuation.ts index d8abf3c5a5..2a736883e0 100644 --- a/scripts/check-tool-turn-continuation.ts +++ b/scripts/check-tool-turn-continuation.ts @@ -20,9 +20,9 @@ interface ModelState { phase: Phase[] // Task.assistantMessageContent partial flag plus tool_result presence stream: "open" | "complete" // Task.didCompleteReadingStream index: number // Task.currentStreamingContentIndex - lock: boolean // Task.presentAssistantMessageLocked + lock: boolean // Task.dispatchState !== DispatchState.IDLE owner: number - pending: boolean // Task.presentAssistantMessageHasPendingUpdates + pending: boolean // Task.pendingWork !== null latch: boolean // Task.userMessageContentReady abort: boolean threw: boolean diff --git a/src/core/assistant-message/__tests__/presentAssistantMessage-auto-deny.spec.ts b/src/core/assistant-message/__tests__/presentAssistantMessage-auto-deny.spec.ts index 72c343c1eb..c8ab1c8ad5 100644 --- a/src/core/assistant-message/__tests__/presentAssistantMessage-auto-deny.spec.ts +++ b/src/core/assistant-message/__tests__/presentAssistantMessage-auto-deny.spec.ts @@ -1,5 +1,6 @@ // npx vitest run core/assistant-message/__tests__/presentAssistantMessage-auto-deny.spec.ts +import { DispatchState, type PendingWork } from "../../task/dispatchState" import type { Anthropic } from "@anthropic-ai/sdk" import { describe, it, expect, beforeEach, vi } from "vitest" import { presentAssistantMessage } from "../presentAssistantMessage" @@ -69,8 +70,8 @@ interface MockTask { taskId: string instanceId: string abort: boolean - presentAssistantMessageLocked: boolean - presentAssistantMessageHasPendingUpdates: boolean + dispatchState: DispatchState + pendingWork: PendingWork | null currentStreamingContentIndex: number assistantMessageContent: unknown[] userMessageContent: Anthropic.ToolResultBlockParam[] @@ -102,8 +103,8 @@ function buildMockTask(): MockTask { taskId: "test-task-id", instanceId: "test-instance", abort: false, - presentAssistantMessageLocked: false, - presentAssistantMessageHasPendingUpdates: false, + dispatchState: DispatchState.IDLE, + pendingWork: null, currentStreamingContentIndex: 0, assistantMessageContent: [], userMessageContent: [], diff --git a/src/core/assistant-message/__tests__/presentAssistantMessage-custom-tool.spec.ts b/src/core/assistant-message/__tests__/presentAssistantMessage-custom-tool.spec.ts index 081e27e687..ba238ceb7d 100644 --- a/src/core/assistant-message/__tests__/presentAssistantMessage-custom-tool.spec.ts +++ b/src/core/assistant-message/__tests__/presentAssistantMessage-custom-tool.spec.ts @@ -1,5 +1,6 @@ // npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-custom-tool.spec.ts +import { DispatchState } from "../../task/dispatchState" import { describe, it, expect, beforeEach, vi } from "vitest" import { presentAssistantMessage } from "../presentAssistantMessage" import { validateToolUse } from "../../tools/validateToolUse" @@ -67,8 +68,8 @@ describe("presentAssistantMessage - Custom Tool Recording", () => { taskId: "test-task-id", instanceId: "test-instance", abort: false, - presentAssistantMessageLocked: false, - presentAssistantMessageHasPendingUpdates: false, + dispatchState: DispatchState.IDLE, + pendingWork: null, currentStreamingContentIndex: 0, assistantMessageContent: [], userMessageContent: [], @@ -224,17 +225,17 @@ describe("presentAssistantMessage - Custom Tool Recording", () => { }) }) - describe("Presentation lock ownership", () => { + describe("Dispatch state ownership", () => { it("drains an update queued during the helper-to-wrapper handoff", async () => { mockTask.assistantMessageContent = [{ type: "text", content: "initial", partial: true }] let queuedCall: Promise | undefined - let lockedDuringHandoff = false + let busyDuringHandoff = false mockTask.say.mockImplementationOnce(async () => { // The first microtask precedes the helper's await continuation. // The second follows that continuation but precedes the wrapper's. queueMicrotask(() => { queueMicrotask(() => { - lockedDuringHandoff = mockTask.presentAssistantMessageLocked + busyDuringHandoff = mockTask.dispatchState !== DispatchState.IDLE mockTask.assistantMessageContent[0] = { type: "text", content: "final", partial: false } mockTask.didCompleteReadingStream = true queuedCall = presentAssistantMessage(mockTask, baseSnapshot) @@ -245,15 +246,15 @@ describe("presentAssistantMessage - Custom Tool Recording", () => { await presentAssistantMessage(mockTask, baseSnapshot) await queuedCall - expect(lockedDuringHandoff).toBe(true) + expect(busyDuringHandoff).toBe(true) expect(mockTask.say.mock.calls).toEqual([ ["text", "initial", undefined, true], ["text", "final", undefined, false], ]) expect(mockTask.currentStreamingContentIndex).toBe(1) expect(mockTask.userMessageContentReady).toBe(true) - expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(false) - expect(mockTask.presentAssistantMessageLocked).toBe(false) + expect(mockTask.pendingWork).toBeNull() + expect(mockTask.dispatchState).toBe(DispatchState.IDLE) }) it("stops draining an update queued during the handoff once the task aborts", async () => { @@ -279,7 +280,7 @@ describe("presentAssistantMessage - Custom Tool Recording", () => { mockTask.assistantMessageContent[0] = { type: "text", content: "final", partial: false } mockTask.didCompleteReadingStream = true queuedCall = presentAssistantMessage(mockTask, baseSnapshot) - pendingDuringHandoff = mockTask.presentAssistantMessageHasPendingUpdates + pendingDuringHandoff = mockTask.pendingWork !== null aborted = true }) }) @@ -291,30 +292,30 @@ describe("presentAssistantMessage - Custom Tool Recording", () => { expect(pendingDuringHandoff).toBe(true) expect(mockTask.say).toHaveBeenCalledExactlyOnceWith("text", "initial", undefined, true) expect(mockTask.currentStreamingContentIndex).toBe(0) - expect(mockTask.presentAssistantMessageLocked).toBe(false) + expect(mockTask.dispatchState).toBe(DispatchState.IDLE) }) - it("holds one lock across consecutive blocks and queues overlapping calls", async () => { + it("holds SERIAL across consecutive blocks and queues overlapping calls", async () => { mockTask.assistantMessageContent = [ { type: "text", content: "first", partial: false }, { type: "text", content: "second", partial: false }, ] mockTask.didCompleteReadingStream = true - const lockChanges: boolean[] = [] - let locked = false - Object.defineProperty(mockTask, "presentAssistantMessageLocked", { + const stateChanges: DispatchState[] = [] + let state = DispatchState.IDLE + Object.defineProperty(mockTask, "dispatchState", { configurable: true, - get: () => locked, - set: (value: boolean) => { - locked = value - lockChanges.push(value) + get: () => state, + set: (value: DispatchState) => { + state = value + stateChanges.push(value) }, }) mockTask.say.mockImplementation(async (_type: string, content: string) => { - expect(mockTask.presentAssistantMessageLocked).toBe(true) + expect(mockTask.dispatchState).toBe(DispatchState.SERIAL) if (content === "second") { await presentAssistantMessage(mockTask, baseSnapshot) - expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(true) + expect(mockTask.pendingWork).toEqual({ mode: "serial" }) } }) @@ -326,8 +327,8 @@ describe("presentAssistantMessage - Custom Tool Recording", () => { ]) expect(mockTask.currentStreamingContentIndex).toBe(2) expect(mockTask.userMessageContentReady).toBe(true) - expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(false) - expect(lockChanges).toEqual([true, false]) + expect(mockTask.pendingWork).toBeNull() + expect(stateChanges).toEqual([DispatchState.SERIAL, DispatchState.IDLE]) }) it("consumes a pending update for a partial block without repeating it indefinitely", async () => { @@ -336,11 +337,11 @@ describe("presentAssistantMessage - Custom Tool Recording", () => { .mockImplementationOnce(async () => { mockTask.assistantMessageContent[0] = { type: "text", content: "updated", partial: true } await presentAssistantMessage(mockTask, baseSnapshot) - expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(true) + expect(mockTask.pendingWork).toEqual({ mode: "serial" }) }) .mockImplementationOnce(async () => { - expect(mockTask.presentAssistantMessageLocked).toBe(true) - expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(false) + expect(mockTask.dispatchState).toBe(DispatchState.SERIAL) + expect(mockTask.pendingWork).toBeNull() }) await presentAssistantMessage(mockTask, baseSnapshot) @@ -350,8 +351,8 @@ describe("presentAssistantMessage - Custom Tool Recording", () => { ["text", "updated", undefined, true], ]) expect(mockTask.currentStreamingContentIndex).toBe(0) - expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(false) - expect(mockTask.presentAssistantMessageLocked).toBe(false) + expect(mockTask.pendingWork).toBeNull() + expect(mockTask.dispatchState).toBe(DispatchState.IDLE) }) it("stops internal continuation when the task aborts between blocks", async () => { @@ -367,10 +368,10 @@ describe("presentAssistantMessage - Custom Tool Recording", () => { expect(mockTask.say).toHaveBeenCalledExactlyOnceWith("text", "first", undefined, false) expect(mockTask.currentStreamingContentIndex).toBe(1) - expect(mockTask.presentAssistantMessageLocked).toBe(false) + expect(mockTask.dispatchState).toBe(DispatchState.IDLE) }) - it("releases the presentation lock when a later block throws", async () => { + it("returns to IDLE when a later block throws", async () => { mockTask.assistantMessageContent = [ { type: "text", content: "first", partial: false }, { type: "text", content: "second", partial: false }, @@ -381,12 +382,12 @@ describe("presentAssistantMessage - Custom Tool Recording", () => { expect(mockTask.say).toHaveBeenCalledTimes(2) expect(mockTask.currentStreamingContentIndex).toBe(1) - expect(mockTask.presentAssistantMessageLocked).toBe(false) + expect(mockTask.dispatchState).toBe(DispatchState.IDLE) }) }) describe("Custom tool error recording", () => { - it("releases the presentation lock when dispatch throws", async () => { + it("returns to IDLE when dispatch throws", async () => { mockTask.assistantMessageContent = [ { type: "tool_use", @@ -404,7 +405,7 @@ describe("presentAssistantMessage - Custom Tool Recording", () => { try { await expect(presentAssistantMessage(mockTask, baseSnapshot)).rejects.toThrow("dispatch failed") - expect(mockTask.presentAssistantMessageLocked).toBe(false) + expect(mockTask.dispatchState).toBe(DispatchState.IDLE) } finally { handleSpy.mockRestore() } diff --git a/src/core/assistant-message/__tests__/presentAssistantMessage-images.spec.ts b/src/core/assistant-message/__tests__/presentAssistantMessage-images.spec.ts index c906cf6b05..70da666bc2 100644 --- a/src/core/assistant-message/__tests__/presentAssistantMessage-images.spec.ts +++ b/src/core/assistant-message/__tests__/presentAssistantMessage-images.spec.ts @@ -1,5 +1,6 @@ // npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-images.spec.ts +import { DispatchState } from "../../task/dispatchState" import { describe, it, expect, beforeEach, vi } from "vitest" import { presentAssistantMessage } from "../presentAssistantMessage" import { validateToolUse } from "../../tools/validateToolUse" @@ -40,8 +41,8 @@ describe("presentAssistantMessage - Image Handling in Native Tool Calling", () = taskId: "test-task-id", instanceId: "test-instance", abort: false, - presentAssistantMessageLocked: false, - presentAssistantMessageHasPendingUpdates: false, + dispatchState: DispatchState.IDLE, + pendingWork: null, currentStreamingContentIndex: 0, assistantMessageContent: [], userMessageContent: [], diff --git a/src/core/assistant-message/__tests__/presentAssistantMessage-tool-repetition.spec.ts b/src/core/assistant-message/__tests__/presentAssistantMessage-tool-repetition.spec.ts index 0121fed928..ccf24f3f58 100644 --- a/src/core/assistant-message/__tests__/presentAssistantMessage-tool-repetition.spec.ts +++ b/src/core/assistant-message/__tests__/presentAssistantMessage-tool-repetition.spec.ts @@ -1,5 +1,6 @@ // npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-tool-repetition.spec.ts +import { DispatchState } from "../../task/dispatchState" import { providerIdentifiers } from "@roo-code/types" import type { ToolUse } from "../../../shared/tools" @@ -153,8 +154,8 @@ describe("presentAssistantMessage - Tool Repetition Detection", () => { taskId: "test-task-id", instanceId: "test-instance", abort: false, - presentAssistantMessageLocked: false, - presentAssistantMessageHasPendingUpdates: false, + dispatchState: DispatchState.IDLE, + pendingWork: null, currentStreamingContentIndex: 0, assistantMessageContent: [], userMessageContent: [], @@ -217,8 +218,8 @@ describe("presentAssistantMessage - Tool Repetition Detection", () => { */ function startNextTurn(blocks: ToolUse[]) { mockTask.currentStreamingContentIndex = 0 - mockTask.presentAssistantMessageLocked = false - mockTask.presentAssistantMessageHasPendingUpdates = false + mockTask.dispatchState = DispatchState.IDLE + mockTask.pendingWork = null mockTask.didAlreadyUseTool = false mockTask.didRejectTool = false mockTask.userMessageContent = [] diff --git a/src/core/assistant-message/__tests__/presentAssistantMessage-tool-usage-attribution.spec.ts b/src/core/assistant-message/__tests__/presentAssistantMessage-tool-usage-attribution.spec.ts index d2e04b89a8..0aee40ef81 100644 --- a/src/core/assistant-message/__tests__/presentAssistantMessage-tool-usage-attribution.spec.ts +++ b/src/core/assistant-message/__tests__/presentAssistantMessage-tool-usage-attribution.spec.ts @@ -1,5 +1,6 @@ // npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-tool-usage-attribution.spec.ts +import { DispatchState, type PendingWork } from "../../task/dispatchState" import type { Anthropic } from "@anthropic-ai/sdk" import { TelemetryEventName } from "@roo-code/types" import { describe, it, expect, beforeEach, vi, type Mock } from "vitest" @@ -60,8 +61,8 @@ interface MockTask { taskId: string instanceId: string abort: boolean - presentAssistantMessageLocked: boolean - presentAssistantMessageHasPendingUpdates: boolean + dispatchState: DispatchState + pendingWork: PendingWork | null currentStreamingContentIndex: number assistantMessageContent: unknown[] userMessageContent: Anthropic.ToolResultBlockParam[] @@ -100,8 +101,8 @@ describe("presentAssistantMessage - tool usage attribution", () => { taskId: "test-task-id", instanceId: "test-instance", abort: false, - presentAssistantMessageLocked: false, - presentAssistantMessageHasPendingUpdates: false, + dispatchState: DispatchState.IDLE, + pendingWork: null, currentStreamingContentIndex: 0, assistantMessageContent: [], userMessageContent: [], diff --git a/src/core/assistant-message/__tests__/presentAssistantMessage-unknown-tool.spec.ts b/src/core/assistant-message/__tests__/presentAssistantMessage-unknown-tool.spec.ts index 157dade953..2a3e72b802 100644 --- a/src/core/assistant-message/__tests__/presentAssistantMessage-unknown-tool.spec.ts +++ b/src/core/assistant-message/__tests__/presentAssistantMessage-unknown-tool.spec.ts @@ -1,5 +1,6 @@ // npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-unknown-tool.spec.ts +import { DispatchState } from "../../task/dispatchState" import { describe, it, expect, beforeEach, vi } from "vitest" import { presentAssistantMessage } from "../presentAssistantMessage" import { isValidToolName, validateToolUse } from "../../tools/validateToolUse" @@ -42,8 +43,8 @@ describe("presentAssistantMessage - Unknown Tool Handling", () => { taskId: "test-task-id", instanceId: "test-instance", abort: false, - presentAssistantMessageLocked: false, - presentAssistantMessageHasPendingUpdates: false, + dispatchState: DispatchState.IDLE, + pendingWork: null, currentStreamingContentIndex: 0, assistantMessageContent: [], userMessageContent: [], @@ -243,7 +244,7 @@ describe("presentAssistantMessage - Unknown Tool Handling", () => { expect(mockTask.say).toHaveBeenCalledTimes(2) expect(mockTask.say).toHaveBeenNthCalledWith(1, "text", "streaming", undefined, true) expect(mockTask.say).toHaveBeenNthCalledWith(2, "text", "updated", undefined, true) - expect(mockTask.presentAssistantMessageHasPendingUpdates).toBe(false) + expect(mockTask.pendingWork).toBeNull() expect(mockTask.currentStreamingContentIndex).toBe(0) }) diff --git a/src/core/assistant-message/__tests__/presentAssistantMessage.spec.ts b/src/core/assistant-message/__tests__/presentAssistantMessage.spec.ts new file mode 100644 index 0000000000..eb9b095bc0 --- /dev/null +++ b/src/core/assistant-message/__tests__/presentAssistantMessage.spec.ts @@ -0,0 +1,108 @@ +// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage.spec.ts + +import { describe, it, expect, beforeEach, vi, type Mock } from "vitest" +import { presentAssistantMessage } from "../presentAssistantMessage" +import { DispatchState, transitionDispatch } from "../../task/dispatchState" +import type { Task } from "../../task/Task" +import type { RequestPolicySnapshot } from "../../prompts/tools/effective-tool-policy" + +vi.mock("../../task/Task") +vi.mock("@roo-code/telemetry", () => ({ + TelemetryService: { + instance: { + captureToolUsage: vi.fn(), + captureConsecutiveMistakeError: vi.fn(), + }, + }, +})) + +const snapshot: RequestPolicySnapshot = { disabledTools: [], customModes: [] } + +describe("presentAssistantMessage dispatch re-entrancy", () => { + let task: Task + let say: Mock + + beforeEach(() => { + say = vi.fn().mockResolvedValue(undefined) + // Double provides only the members the text-block path reads. + task = { + taskId: "task-id", + instanceId: "instance-id", + abort: false, + dispatchState: DispatchState.IDLE, + pendingWork: null, + currentStreamingContentIndex: 0, + assistantMessageContent: [{ type: "text", content: "partial", partial: true }], + userMessageContent: [], + didCompleteReadingStream: false, + activeToolContexts: new Map(), + say, + } as unknown as Task + }) + + it.each([DispatchState.SERIAL, DispatchState.PARALLEL])("queues pending work while %s", async (state) => { + task.dispatchState = state + + await presentAssistantMessage(task, snapshot) + + expect(task.pendingWork).toEqual({ mode: "serial" }) + expect(task.dispatchState).toBe(state) + expect(say).not.toHaveBeenCalled() + }) + + it("drains pending work once after the first call completes", async () => { + say.mockImplementationOnce(async () => { + expect(task.dispatchState).toBe(DispatchState.SERIAL) + await presentAssistantMessage(task, snapshot) + expect(task.pendingWork).toEqual({ mode: "serial" }) + }) + + await presentAssistantMessage(task, snapshot) + + expect(say).toHaveBeenCalledTimes(2) + expect(task.pendingWork).toBeNull() + expect(task.dispatchState).toBe(DispatchState.IDLE) + }) + + it("runs at most twice for N rapid re-entrant calls", async () => { + say.mockImplementationOnce(async () => { + await Promise.all(Array.from({ length: 10 }, () => presentAssistantMessage(task, snapshot))) + }) + + await presentAssistantMessage(task, snapshot) + + expect(say).toHaveBeenCalledTimes(2) + expect(task.dispatchState).toBe(DispatchState.IDLE) + }) + + it("returns to IDLE when the presenter throws", async () => { + say.mockRejectedValueOnce(new Error("boom")) + + await expect(presentAssistantMessage(task, snapshot)).rejects.toThrow("boom") + + expect(task.dispatchState).toBe(DispatchState.IDLE) + }) +}) + +describe("transitionDispatch", () => { + const { IDLE, SERIAL, PARALLEL } = DispatchState + + it.each([ + [IDLE, SERIAL], + [IDLE, PARALLEL], + [SERIAL, IDLE], + [PARALLEL, IDLE], + ])("accepts %s -> %s", (from, to) => { + expect(transitionDispatch(from, to)).toBe(to) + }) + + it.each([ + [IDLE, IDLE], + [SERIAL, SERIAL], + [SERIAL, PARALLEL], + [PARALLEL, SERIAL], + [PARALLEL, PARALLEL], + ])("rejects %s -> %s", (from, to) => { + expect(() => transitionDispatch(from, to)).toThrow("invalid transition") + }) +}) diff --git a/src/core/assistant-message/presentAssistantMessage.ts b/src/core/assistant-message/presentAssistantMessage.ts index fa9057d6d5..73d48b2e7c 100644 --- a/src/core/assistant-message/presentAssistantMessage.ts +++ b/src/core/assistant-message/presentAssistantMessage.ts @@ -1,4 +1,5 @@ import { markToolRejected } from "../task/toolExecutionContext" +import { DispatchState, transitionDispatch } from "../task/dispatchState" import { serializeError } from "serialize-error" import { Anthropic } from "@anthropic-ai/sdk" @@ -102,12 +103,12 @@ export async function presentAssistantMessage(cline: Task, policySnapshot: Reque ) } - if (cline.presentAssistantMessageLocked) { - cline.presentAssistantMessageHasPendingUpdates = true + if (cline.dispatchState !== DispatchState.IDLE) { + cline.pendingWork = { mode: "serial" } return } - cline.presentAssistantMessageLocked = true + cline.dispatchState = transitionDispatch(cline.dispatchState, DispatchState.SERIAL) try { // Drain updates queued while the lock was held, including one that // arrives after the helper's final check but before this continuation. @@ -115,11 +116,11 @@ export async function presentAssistantMessage(cline: Task, policySnapshot: Reque // synchronous step, so no update can be stranded behind the lock. do { await presentAssistantMessageBlock(cline, policySnapshot) - } while (!cline.abort && cline.presentAssistantMessageHasPendingUpdates) + } while (!cline.abort && cline.pendingWork) } finally { // Tool handlers and provider-state reads can reject. Never strand the // task behind a dispatch lock after the presenter has unwound. - cline.presentAssistantMessageLocked = false + cline.dispatchState = DispatchState.IDLE } } @@ -130,7 +131,7 @@ async function presentAssistantMessageBlock(cline: Task, policySnapshot: Request // Each internal pass consumes the pending update. New updates arriving // during an awaited operation can request another pass. - cline.presentAssistantMessageHasPendingUpdates = false + cline.pendingWork = null if (cline.currentStreamingContentIndex >= cline.assistantMessageContent.length) { // This may happen if the last content block was completed before diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index ebb362f7eb..24f406cbc1 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -109,6 +109,7 @@ import { FileContextTracker } from "../context-tracking/FileContextTracker" import { RooIgnoreController } from "../ignore/RooIgnoreController" import { RooProtectedController } from "../protect/RooProtectedController" import { type AssistantMessageContent, presentAssistantMessage } from "../assistant-message" +import { DispatchState, type PendingWork } from "./dispatchState" import { type RequestPolicySnapshot } from "../prompts/tools/effective-tool-policy" import { NativeToolCallParser } from "../assistant-message/NativeToolCallParser" import { manageContext, willManageContext } from "../context-management" @@ -516,8 +517,8 @@ export class Task extends EventEmitter implements TaskLike { currentStreamingContentIndex = 0 currentStreamingDidCheckpoint = false assistantMessageContent: AssistantMessageContent[] = [] - presentAssistantMessageLocked = false - presentAssistantMessageHasPendingUpdates = false + dispatchState: DispatchState = DispatchState.IDLE + pendingWork: PendingWork | null = null userMessageContent: (Anthropic.TextBlockParam | Anthropic.ImageBlockParam | Anthropic.ToolResultBlockParam)[] = [] userMessageContentReady = false @@ -609,7 +610,7 @@ export class Task extends EventEmitter implements TaskLike { // A non-abort presenter rejection releases the lock and is logged, but // nothing schedules another presentation pass. Gating on the index would // strand this wait after exactly that failure. - if (!this.didCompleteReadingStream || this.presentAssistantMessageLocked) { + if (!this.didCompleteReadingStream || this.dispatchState !== DispatchState.IDLE) { return false } @@ -4135,8 +4136,8 @@ export class Task extends EventEmitter implements TaskLike { // for the turn that earned it: a message the user queues afterward // (a new turn) is fair game for the next ask to answer. this.blanketDeniedCommandThisTurn = false - this.presentAssistantMessageLocked = false - this.presentAssistantMessageHasPendingUpdates = false + this.dispatchState = DispatchState.IDLE + this.pendingWork = null // No legacy text-stream tool parser. this.streamingToolCallIndices.clear() const nativeToolCallParserScope = NativeToolCallParser.createScope() diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index 3066613263..1df6ee3ca7 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -1,6 +1,7 @@ import { WebviewFocusTracker } from "../../webview/WebviewFocusTracker" // npx vitest core/task/__tests__/Task.spec.ts +import { DispatchState } from "../dispatchState" import * as fsReal from "fs" import * as os from "os" import * as path from "path" @@ -646,9 +647,11 @@ describe("Cline", () => { expect(readiness()).toBe(false) task.didCompleteReadingStream = true expect(readiness()).toBe(true) - task.presentAssistantMessageLocked = true + task.dispatchState = DispatchState.SERIAL expect(readiness()).toBe(false) - task.presentAssistantMessageLocked = false + task.dispatchState = DispatchState.PARALLEL + expect(readiness()).toBe(false) + task.dispatchState = DispatchState.IDLE expect(readiness()).toBe(true) task.assistantMessageContent[0].partial = true expect(readiness()).toBe(false) @@ -694,7 +697,7 @@ describe("Cline", () => { startTask: false, }) task.didCompleteReadingStream = true - task.presentAssistantMessageLocked = false + task.dispatchState = DispatchState.IDLE task.assistantMessageContent = [ { type: "tool_use", @@ -915,7 +918,7 @@ describe("Cline", () => { content: "File: package.json\nfinished", }), ]) - expect(task.presentAssistantMessageLocked).toBe(false) + expect(task.dispatchState).toBe(DispatchState.IDLE) expect(sawCompleteToolTurn).toBe(true) } finally { vi.mocked(pWaitFor).mockImplementation(async () => {}) @@ -982,13 +985,13 @@ describe("Cline", () => { await resultPushed // Give the request loop several poll intervals to continue early. await new Promise((resolve) => setTimeout(resolve, 50)) - expect(task.presentAssistantMessageLocked).toBe(true) + expect(task.dispatchState).toBe(DispatchState.SERIAL) expect(attemptApiRequestSpy).toHaveBeenCalledTimes(1) releaseHandler?.() await run expect(attemptApiRequestSpy).toHaveBeenCalledTimes(2) - expect(task.presentAssistantMessageLocked).toBe(false) + expect(task.dispatchState).toBe(DispatchState.IDLE) } finally { releaseHandler?.() vi.mocked(pWaitFor).mockImplementation(async () => {}) @@ -1056,7 +1059,7 @@ describe("Cline", () => { expect(attemptApiRequestSpy).toHaveBeenCalledOnce() expect(sawReadyWhileAborted).toBe(true) - expect(task.presentAssistantMessageLocked).toBe(false) + expect(task.dispatchState).toBe(DispatchState.IDLE) } finally { vi.mocked(pWaitFor).mockImplementation(async () => {}) readinessSpy.mockRestore() diff --git a/src/core/task/__tests__/ToolExecutionContext.spec.ts b/src/core/task/__tests__/ToolExecutionContext.spec.ts index f466c05596..649a7af82f 100644 --- a/src/core/task/__tests__/ToolExecutionContext.spec.ts +++ b/src/core/task/__tests__/ToolExecutionContext.spec.ts @@ -1,5 +1,6 @@ // npx vitest run core/task/__tests__/ToolExecutionContext.spec.ts +import { DispatchState } from "../dispatchState" import type { Anthropic } from "@anthropic-ai/sdk" import { describe, it, expect, beforeEach, vi } from "vitest" import { presentAssistantMessage } from "../../assistant-message/presentAssistantMessage" @@ -44,8 +45,8 @@ function buildMockTask() { taskId: "test-task-id", instanceId: "test-instance", abort: false, - presentAssistantMessageLocked: false, - presentAssistantMessageHasPendingUpdates: false, + dispatchState: DispatchState.IDLE, + pendingWork: null, currentStreamingContentIndex: 0, assistantMessageContent: [] as unknown[], userMessageContent: [] as Anthropic.ToolResultBlockParam[], @@ -158,7 +159,7 @@ describe("ToolExecutionContext", () => { mockTask.ask.mockResolvedValueOnce({ response: "noButtonClicked" }) await presentAssistantMessage(asTask(mockTask), baseSnapshot) - mockTask.presentAssistantMessageLocked = false + mockTask.dispatchState = DispatchState.IDLE mockTask.currentStreamingContentIndex = 1 await presentAssistantMessage(asTask(mockTask), baseSnapshot) diff --git a/src/core/task/dispatchState.ts b/src/core/task/dispatchState.ts new file mode 100644 index 0000000000..c55bb9ce79 --- /dev/null +++ b/src/core/task/dispatchState.ts @@ -0,0 +1,20 @@ +export enum DispatchState { + IDLE = "idle", + SERIAL = "serial", + PARALLEL = "parallel", +} + +export type PendingWork = { mode: "serial" | "parallel" } + +const VALID_TRANSITIONS: Record = { + [DispatchState.IDLE]: [DispatchState.SERIAL, DispatchState.PARALLEL], + [DispatchState.SERIAL]: [DispatchState.IDLE], + [DispatchState.PARALLEL]: [DispatchState.IDLE], +} + +export function transitionDispatch(from: DispatchState, to: DispatchState): DispatchState { + if (!VALID_TRANSITIONS[from].includes(to)) { + throw new Error(`[DispatchState] invalid transition ${from} -> ${to}`) + } + return to +}