feat(write-to-file): per-task partial stream state + cleanup primitives (split 2/6 of #1066)

src/__tests__/removeClineFromStack-delegation.spec.ts

@@ -5,6 +5,7 @@ import { ClineProvider } from "../core/webview/ClineProvider"
 import { TaskRegistry } from "../core/task/TaskRegistry"
 import { PendingActionSettlementError, type Task } from "../core/task/Task"
 import { makeProviderStub } from "./helpers/provider-stub"
+import { writeToFileTool } from "../core/tools/WriteToFileTool"
 
 type MockTask = Pick<Task, "taskId" | "instanceId"> &
 	Partial<Pick<Task, "parentTaskId" | "abort" | "abandoned">> & {
@@ -208,6 +209,46 @@ describe("ClineProvider failed history restoration cleanup", () => {
 		expect(task.dispose).toHaveBeenCalledOnce()
 	})
 
+	it("releases the tool's per-task state before a directly disposed task loses its listeners", async () => {
+		const key = "failed-history-task.inst-1"
+		const order: string[] = []
+		const task = {
+			taskId: "failed-history-task",
+			instanceId: "inst-1",
+			emit: vi.fn(),
+			once: vi.fn(),
+			off: vi.fn(),
+			dispose: vi.fn().mockImplementation(() => {
+				// Dispose removes every listener, so this is the last moment the abort
+				// cleanup could still have run; the entry must already be gone here.
+				order.push(writeToFileTool["taskPartialStreamState"].has(key) ? "state-retained" : "state-cleared")
+				order.push("dispose")
+				return Promise.resolve()
+			}),
+		} as unknown as Task
+		// Fixture: the state entry the tool would have created during a partial stream.
+		writeToFileTool["taskPartialStreamState"].set(key, {
+			lastSeenPartialPath: undefined,
+			streamFailed: false,
+			streamError: undefined,
+			task,
+			abortCleanup: () => {},
+		})
+		const taskEventListeners = new Map([[task, [vi.fn()]]])
+		const taskRegistry = new TaskRegistry()
+		taskRegistry.push(task)
+		const provider = { taskRegistry, taskEventListeners, log: vi.fn() } as unknown as ClineProvider
+
+		await privateClineProvider.cleanupFailedHistoryTask.call(
+			provider,
+			task,
+			new PendingActionSettlementError("settlement failed"),
+		)
+
+		expect(order).toEqual(["state-cleared", "dispose"])
+		expect(writeToFileTool["taskPartialStreamState"].size).toBe(0)
+	})
+
 	it("keeps the task active after an unrelated history resume failure", async () => {
 		const cleanupListener = vi.fn()
 		const task = {

src/core/assistant-message/__tests__/presentAssistantMessage-validation-rejection.spec.ts

@@ -0,0 +1,260 @@
+// npx vitest run core/assistant-message/__tests__/presentAssistantMessage-validation-rejection.spec.ts
+
+import { describe, it, expect, beforeEach, vi } from "vitest"
+
+import { providerIdentifiers } from "@roo-code/types/provider-identifiers"
+
+import { presentAssistantMessage } from "../presentAssistantMessage"
+import type { Task } from "../../task/Task"
+
+const mockRelease = vi.hoisted(() => vi.fn().mockResolvedValue(undefined))
+const mockHandle = vi.hoisted(() => vi.fn().mockResolvedValue(undefined))
+const mockValidate = vi.hoisted(() => vi.fn())
+
+vi.mock("../../task/Task")
+vi.mock("../../tools/validateToolUse", () => ({
+	validateToolUse: mockValidate,
+	isValidToolName: vi.fn(() => true),
+}))
+vi.mock("../../tools/NewTaskTool", () => ({ newTaskTool: { handle: vi.fn() } }))
+vi.mock("../../tools/WriteToFileTool", () => ({
+	writeToFileTool: { handle: mockHandle, releaseStreamAfterValidationRejection: mockRelease },
+}))
+vi.mock("@roo-code/telemetry", () => ({
+	TelemetryService: {
+		instance: {
+			captureToolUsage: vi.fn(),
+			captureConsecutiveMistakeError: vi.fn(),
+			captureException: vi.fn(),
+		},
+	},
+}))
+
+interface ToolResultBlock {
+	type: string
+	tool_use_id?: string
+	content?: string
+	is_error?: boolean
+}
+
+/**
+ * Structural double for the presenter's Task surface. The presenter reads far more of Task
+ * than the two rejection branches touch, so the double carries only what this file drives;
+ * it is handed over through unknown (the repo's pattern for presenter-level doubles, see
+ * presentAssistantMessage-tool-usage-attribution.spec.ts) rather than any.
+ */
+interface PresenterTask {
+	taskId: string
+	instanceId: string
+	abort: boolean
+	presentAssistantMessageLocked: boolean
+	presentAssistantMessageHasPendingUpdates: boolean
+	currentStreamingContentIndex: number
+	assistantMessageContent: Record<string, unknown>[]
+	userMessageContent: ToolResultBlock[]
+	didCompleteReadingStream: boolean
+	didRejectTool: boolean
+	didAlreadyUseTool: boolean
+	consecutiveMistakeCount: number
+	consecutiveMistakeLimit: number
+	apiConfiguration: { apiProvider: string }
+	clineMessages: unknown[]
+	getTaskMode: ReturnType<typeof vi.fn>
+	api: { getModel: () => { id: string; info: Record<string, unknown> } }
+	recordToolUsage: ReturnType<typeof vi.fn>
+	recordToolError: ReturnType<typeof vi.fn>
+	toolRepetitionDetector: { check: ReturnType<typeof vi.fn> }
+	providerRef: { deref: () => { getState: ReturnType<typeof vi.fn> } }
+	say: ReturnType<typeof vi.fn>
+	ask: ReturnType<typeof vi.fn>
+	pushToolResultToUserContent: ReturnType<typeof vi.fn>
+}
+
+describe("presentAssistantMessage - a rejected write_to_file releases its stream", () => {
+	let mockTask: PresenterTask
+
+	beforeEach(() => {
+		vi.clearAllMocks()
+		mockRelease.mockResolvedValue(undefined)
+		mockHandle.mockResolvedValue(undefined)
+		mockValidate.mockImplementation(() => {
+			throw new Error("write_to_file is not allowed in this mode")
+		})
+		mockTask = {
+			taskId: "validation-task",
+			instanceId: "inst-1",
+			abort: false,
+			presentAssistantMessageLocked: false,
+			presentAssistantMessageHasPendingUpdates: false,
+			currentStreamingContentIndex: 0,
+			assistantMessageContent: [],
+			userMessageContent: [],
+			didCompleteReadingStream: false,
+			didRejectTool: false,
+			didAlreadyUseTool: false,
+			consecutiveMistakeCount: 0,
+			consecutiveMistakeLimit: 3,
+			apiConfiguration: { apiProvider: providerIdentifiers.anthropic },
+			clineMessages: [],
+			getTaskMode: vi.fn().mockResolvedValue("code"),
+			api: { getModel: () => ({ id: "test-model", info: {} }) },
+			recordToolUsage: vi.fn(),
+			recordToolError: vi.fn(),
+			toolRepetitionDetector: { check: vi.fn().mockReturnValue({ allowExecution: true }) },
+			providerRef: { deref: () => ({ getState: vi.fn().mockResolvedValue({ mode: "code", customModes: [] }) }) },
+			say: vi.fn().mockResolvedValue(undefined),
+			ask: vi.fn().mockResolvedValue({ response: "yesButtonClicked" }),
+
+			pushToolResultToUserContent: vi.fn().mockImplementation((toolResult) => {
+				mockTask.userMessageContent.push(toolResult)
+				return true
+			}),
+		}
+	})
+
+	const writeBlock = () => [
+		{
+			type: "tool_use",
+			id: "call-write-1",
+			name: "write_to_file",
+			params: { path: "a.ts", content: "partial from the stream" },
+			// The presenter breaks earlier for a known tool whose nativeArgs never arrived, so the
+			// block has to carry them to reach the validation step.
+			nativeArgs: { path: "a.ts", content: "partial from the stream" },
+			partial: false,
+		},
+	]
+
+	it("releases the streamed state when validation rejects the completed block", async () => {
+		// A partial delta is never validated, so streaming can already have registered this
+		// task's write_to_file state and opened a preview by the time the completed block is
+		// rejected. The loop breaks before writeToFileTool.handle() runs, so nothing else
+		// releases them and the task carries a stale stream into its next write.
+		mockTask.assistantMessageContent = writeBlock()
+
+		await presentAssistantMessage(mockTask as unknown as Task)
+
+		expect(mockRelease).toHaveBeenCalledTimes(1)
+		expect(mockRelease).toHaveBeenCalledWith(mockTask)
+		expect(mockHandle).not.toHaveBeenCalled()
+		// The validation error stays the tool result the model sees.
+		const toolResult = mockTask.userMessageContent.find((item) => item.type === "tool_result")
+		if (!toolResult) {
+			throw new Error("expected a tool_result for the rejected call")
+		}
+		expect(toolResult.is_error).toBe(true)
+		expect(toolResult.content).toContain("not allowed in this mode")
+	})
+
+	it("releases the streamed state when the repetition guard refuses the completed block", async () => {
+		// Same family: the block is refused before handle() runs, so the stream owns nobody
+		// but this branch.
+		mockValidate.mockReturnValue(undefined)
+		mockTask.toolRepetitionDetector.check = vi.fn().mockReturnValue({
+			allowExecution: false,
+			askUser: { messageKey: "mistake_limit_reached", messageDetail: "repeated" },
+		})
+		mockTask.assistantMessageContent = writeBlock()
+
+		await presentAssistantMessage(mockTask as unknown as Task)
+
+		expect(mockRelease).toHaveBeenCalledTimes(1)
+		expect(mockHandle).not.toHaveBeenCalled()
+		// pushToolResult is the presenter's own closure; what the model receives is the
+		// tool_result in the user content, which is what must still carry the refusal.
+		const toolResult = mockTask.userMessageContent.find((item) => item.type === "tool_result")
+		if (!toolResult) {
+			throw new Error("expected a tool_result for the rejected call")
+		}
+		expect(toolResult.content).toContain("repetition limit reached")
+	})
+
+	it("releases the streamed state when the completed block carries no native arguments", async () => {
+		// A third way to break out of the loop before handle(): the parser never finished the
+		// call. Streaming is not gated by it either, so the same state and preview can exist.
+		mockValidate.mockReturnValue(undefined)
+		mockTask.assistantMessageContent = [
+			{
+				type: "tool_use",
+				id: "call-write-2",
+				name: "write_to_file",
+				params: { path: "a.ts" },
+				partial: false,
+			},
+		]
+
+		await presentAssistantMessage(mockTask as unknown as Task)
+
+		expect(mockRelease).toHaveBeenCalledTimes(1)
+		expect(mockRelease).toHaveBeenCalledWith(mockTask)
+		expect(mockHandle).not.toHaveBeenCalled()
+		const toolResult = mockTask.userMessageContent.find((item) => item.type === "tool_result")
+		if (!toolResult) {
+			throw new Error("expected a tool_result for the rejected call")
+		}
+		expect(toolResult.content).toContain("missing nativeArgs")
+	})
+
+	it("does not release the write_to_file stream for a different tool refused by the repetition guard", async () => {
+		// The release is scoped by tool name in every branch. A repeated read_file must not
+		// reach through to write_to_file's state.
+		mockValidate.mockReturnValue(undefined)
+		mockTask.toolRepetitionDetector.check = vi.fn().mockReturnValue({
+			allowExecution: false,
+			askUser: { messageKey: "mistake_limit_reached", messageDetail: "repeated" },
+		})
+		mockTask.assistantMessageContent = [
+			{
+				type: "tool_use",
+				id: "call-read-2",
+				name: "read_file",
+				params: { path: "a.ts" },
+				nativeArgs: { path: "a.ts" },
+				partial: false,
+			},
+		]
+
+		await presentAssistantMessage(mockTask as unknown as Task)
+
+		expect(mockRelease).not.toHaveBeenCalled()
+		expect(mockHandle).not.toHaveBeenCalled()
+	})
+
+	it("does not release the write_to_file stream for a different tool with no native arguments", async () => {
+		// The name check guards this branch too. Without a negative control here, a mutant
+		// that releases for every tool in this branch survives.
+		mockValidate.mockReturnValue(undefined)
+		mockTask.assistantMessageContent = [
+			{
+				type: "tool_use",
+				id: "call-read-3",
+				name: "read_file",
+				params: { path: "a.ts" },
+				partial: false,
+			},
+		]
+
+		await presentAssistantMessage(mockTask as unknown as Task)
+
+		expect(mockRelease).not.toHaveBeenCalled()
+		expect(mockHandle).not.toHaveBeenCalled()
+	})
+
+	it("leaves the stream alone when a rejected tool never streamed", async () => {
+		// The release is write_to_file scoped: a rejected read_file must not touch it.
+		mockTask.assistantMessageContent = [
+			{
+				type: "tool_use",
+				id: "call-read-1",
+				name: "read_file",
+				params: { path: "a.ts" },
+				nativeArgs: { path: "a.ts" },
+				partial: false,
+			},
+		]
+
+		await presentAssistantMessage(mockTask as unknown as Task)
+
+		expect(mockRelease).not.toHaveBeenCalled()
+	})
+})

src/core/assistant-message/presentAssistantMessage.ts

@@ -572,6 +572,14 @@ async function presentAssistantMessageBlock(cline: Task): Promise<void> {
 						is_error: true,
 					})
 
+					// A partial delta is never validated or parsed, so this task's write_to_file stream
+					// state and its preview can already exist when the completed block turns out to
+					// carry no native arguments. The loop breaks here without reaching handle(), so
+					// nothing else releases them.
+					if (block.name === "write_to_file") {
+						await writeToFileTool.releaseStreamAfterValidationRejection(cline)
+					}
+
 					break
 				}
 			}
@@ -779,6 +787,14 @@ async function presentAssistantMessageBlock(cline: Task): Promise<void> {
 						error.message,
 					)
 
+					// A partial delta may already have registered this task's write_to_file stream
+					// state and opened a preview. Validation rejects the completed block before
+					// writeToFileTool.handle() runs, so nothing else releases them and the task would
+					// carry a stale stream into its next write.
+					if (block.name === "write_to_file") {
+						await writeToFileTool.releaseStreamAfterValidationRejection(cline)
+					}
+
 					break
 				}
 
@@ -846,6 +862,11 @@ async function presentAssistantMessageBlock(cline: Task): Promise<void> {
 							`Tool call repetition limit reached for ${block.name}. Please try a different approach.`,
 						),
 					)
+					// Same family as the validation branch above: the completed block is refused
+					// before handle() runs, so a stream that had already started owns nobody but here.
+					if (block.name === "write_to_file") {
+						await writeToFileTool.releaseStreamAfterValidationRejection(cline)
+					}
 					break
 				}
 			}

src/core/tools/BaseTool.ts

@@ -98,6 +98,19 @@ export abstract class BaseTool<TName extends ToolName> {
 		this.lastSeenPartialPath = undefined
 	}
 
+	/**
+	 * Teardown boundary for the handle() parse-failure path, where execute() never
+	 * runs. Default: there is no per-task streaming state to release, so the generic
+	 * parse error is what the user sees. A tool that keeps per-task stream state may
+	 * release it and restore any diff document a stream opened - returning true
+	 * suppresses the incidental parse error so the failure is reported exactly once.
+	 * Per-task only: a global teardown would clobber another task that is still
+	 * streaming through this singleton.
+	 */
+	protected async releaseStreamStateOnParseFailure(_task: Task, _callbacks: ToolCallbacks): Promise<boolean> {
+		return false
+	}
+
 	/**
 	 * Main entry point for tool execution.
 	 *
@@ -157,7 +170,14 @@ export abstract class BaseTool<TName extends ToolName> {
 		} catch (error) {
 			console.error(`Error parsing parameters:`, error)
 			const errorMessage = `Failed to parse ${this.name} parameters: ${error instanceof Error ? error.message : String(error)}`
-			await callbacks.handleError(`parsing ${this.name} args`, new Error(errorMessage))
+			// execute() never runs on this path, so a tool that keeps per-task streaming state
+			// must still release THIS task's state and restore any diff document the stream
+			// opened. When the tool reports a more specific failure, the incidental parse error
+			// is suppressed so the failure surfaces exactly once.
+			const reportedStreamFailure = await this.releaseStreamStateOnParseFailure(task, callbacks)
+			if (!reportedStreamFailure) {
+				await callbacks.handleError(`parsing ${this.name} args`, new Error(errorMessage))
+			}
 			// Note: handleError already emits a tool_result via formatResponse.toolError in the caller.
 			// Do NOT call pushToolResult here to avoid duplicate tool_result payloads.
 			return