fix(write-to-file): clean partial state on missing-param and rooignore denial + integrate the #1066 split series (6/6 + FINAL)

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 PrivateClineProviderMethods = {
 	removeClineFromStack: (this: ClineProvider) => ReturnType<ClineProvider["removeClineFromStack"]>
@@ -202,6 +203,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-missing-native-args.spec.ts

@@ -0,0 +1,197 @@
+// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-missing-native-args.spec.ts
+
+import { describe, it, expect, beforeEach, vi } from "vitest"
+import { presentAssistantMessage } from "../presentAssistantMessage"
+import { DispatchState } from "../../task/dispatchState"
+import { type RequestPolicySnapshot } from "../../prompts/tools/effective-tool-policy"
+import { isValidToolName, validateToolUse } from "../../tools/validateToolUse"
+
+// Tool policy is frozen at request build and the presenter refuses to run without the
+// snapshot; these tests only exercise the presenter's bypass exits, so the empty policy is
+// enough (validateToolUse itself is mocked).
+const baseSnapshot: RequestPolicySnapshot = {
+	disabledTools: [],
+	customModes: [],
+}
+
+const mockTeardown = vi.hoisted(() => vi.fn())
+const mockWriteHandle = vi.hoisted(() => vi.fn())
+
+vi.mock("../../task/Task")
+vi.mock("../../tools/validateToolUse", () => ({
+	validateToolUse: vi.fn(),
+	isValidToolName: vi.fn(() => true),
+}))
+vi.mock("../../tools/WriteToFileTool", () => ({
+	writeToFileTool: { handle: mockWriteHandle, teardownAbandonedStream: mockTeardown },
+}))
+vi.mock("@roo-code/telemetry", () => ({
+	TelemetryService: {
+		instance: {
+			captureToolUsage: vi.fn(),
+			captureConsecutiveMistakeError: vi.fn(),
+			captureException: vi.fn(),
+		},
+	},
+}))
+
+type ToolResultBlock = { type: string; tool_use_id: string; content: string; is_error: boolean }
+
+describe("presentAssistantMessage - finalized block without nativeArgs", () => {
+	// The presenter reads only a subset of Task. A Record keeps the double honest without
+	// an `any`; the single cast at the call site is the documented boundary.
+	let mockTask: Record<string, unknown>
+	let userMessageContent: ToolResultBlock[]
+
+	beforeEach(() => {
+		mockTeardown.mockReset()
+		mockWriteHandle.mockReset()
+		userMessageContent = []
+		mockTask = {
+			taskId: "test-task-id",
+			instanceId: "test-instance",
+			abort: false,
+			// main renamed the dispatcher fields (src/core/task/dispatchState.ts): the presenter reads
+			// dispatchState and pendingWork, not the old lock/flag pair.
+			dispatchState: DispatchState.IDLE,
+			pendingWork: null,
+			currentStreamingContentIndex: 0,
+			assistantMessageContent: [
+				{
+					type: "tool_use",
+					name: "write_to_file",
+					params: {},
+					partial: false,
+					// Streaming JSON never parsed: Task completes the block with no nativeArgs.
+					id: "toolu_1",
+					nativeArgs: undefined,
+				},
+			],
+			userMessageContent,
+			didCompleteReadingStream: false,
+			// The presenter keeps per-tool-call context in this map (added on main since this
+			// spec was written); the bypass exits below run after its first use.
+			activeToolContexts: new Map(),
+			didRejectTool: false,
+			didAlreadyUseTool: false,
+			consecutiveMistakeCount: 0,
+			consecutiveMistakeLimit: 3,
+			apiConfiguration: { apiProvider: "test-provider" },
+			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((toolResult: ToolResultBlock) => {
+				userMessageContent.push(toolResult)
+				return true
+			}),
+		}
+	})
+
+	it("tears the write_to_file stream state down instead of dispatching the tool", async () => {
+		await presentAssistantMessage(
+			mockTask as unknown as Parameters<typeof presentAssistantMessage>[0],
+			baseSnapshot,
+		)
+
+		// The malformed call must not be executed, and exactly one error tool_result is
+		// emitted for the provider.
+		expect(isValidToolName).toHaveBeenCalled()
+		expect(mockWriteHandle).not.toHaveBeenCalled()
+		expect(mockTask.pushToolResultToUserContent).toHaveBeenCalledTimes(1)
+		expect(userMessageContent).toEqual([
+			expect.objectContaining({
+				type: "tool_result",
+				tool_use_id: "toolu_1",
+				is_error: true,
+				content: expect.stringContaining("missing nativeArgs"),
+			}),
+		])
+		// The guard bypasses handle(), so the per-task stream state has to be released here:
+		// otherwise the entry, its TaskAborted listener and any streamed diff view leak into
+		// the next API request of the same task.
+		expect(mockTeardown).toHaveBeenCalledTimes(1)
+		expect(mockTeardown).toHaveBeenCalledWith(mockTask)
+	})
+
+	it("tears the stream state down when tool validation rejects the call", async () => {
+		// A mode file restriction (or any validateToolUse failure) lands here. The block
+		// streamed partial deltas - handlePartial ran - and this exit bypasses handle(), so
+		// without the teardown the per-task stream state, its TaskAborted listener and any
+		// streamed diff view leak into the next API request.
+		mockTask.assistantMessageContent = [
+			{
+				type: "tool_use",
+				name: "write_to_file",
+				params: { path: "restricted.ts", content: "partial model output" },
+				partial: false,
+				id: "toolu_2",
+				nativeArgs: { path: "restricted.ts", content: "partial model output" },
+			},
+		]
+		vi.mocked(validateToolUse).mockImplementationOnce(() => {
+			throw new Error("write_to_file is not allowed to write restricted.ts in this mode")
+		})
+
+		await presentAssistantMessage(
+			mockTask as unknown as Parameters<typeof presentAssistantMessage>[0],
+			baseSnapshot,
+		)
+
+		expect(mockWriteHandle).not.toHaveBeenCalled()
+		expect(userMessageContent).toEqual([
+			expect.objectContaining({
+				type: "tool_result",
+				tool_use_id: "toolu_2",
+				is_error: true,
+				content: expect.stringContaining("restricted.ts"),
+			}),
+		])
+		expect(mockTeardown).toHaveBeenCalledTimes(1)
+		expect(mockTeardown).toHaveBeenCalledWith(mockTask)
+	})
+
+	it("tears the stream state down when the tool repetition limit stops the call", async () => {
+		mockTask.assistantMessageContent = [
+			{
+				type: "tool_use",
+				name: "write_to_file",
+				params: { path: "same.ts", content: "same content" },
+				partial: false,
+				id: "toolu_3",
+				nativeArgs: { path: "same.ts", content: "same content" },
+			},
+		]
+		// The detector's result is a tagged union now (allow / soft_block / hard_block): only the
+		// hard block asks the user and stops the call, so that is the exit under test.
+		vi.mocked(mockTask.toolRepetitionDetector as unknown as { check: unknown }).check = vi.fn().mockReturnValue({
+			action: "hard_block",
+			askUser: { messageKey: "tool_repetition", messageDetail: "write_to_file" },
+		})
+
+		await presentAssistantMessage(
+			mockTask as unknown as Parameters<typeof presentAssistantMessage>[0],
+			baseSnapshot,
+		)
+
+		expect(mockWriteHandle).not.toHaveBeenCalled()
+		// The repetition exit reports through the local pushToolResult, which wraps the
+		// message in the tool-error envelope rather than the provider is_error flag.
+		expect(userMessageContent).toEqual([
+			expect.objectContaining({
+				type: "tool_result",
+				tool_use_id: "toolu_3",
+				content: expect.stringContaining("repetition limit reached"),
+			}),
+		])
+		expect(mockTeardown).toHaveBeenCalledTimes(1)
+		expect(mockTeardown).toHaveBeenCalledWith(mockTask)
+	})
+})

src/core/assistant-message/presentAssistantMessage.ts

@@ -586,6 +586,14 @@
 						is_error: true,
 					})
 
+					// A write_to_file call that streamed partial deltas leaves per-task stream state
+					// (and possibly an open diff view) behind, and this guard bypasses handle(), so
+					// the teardown that execute()/onParameterParseFailure would have run never does.
+					// Release it here instead of leaking it into the next API request.
+					if (block.name === "write_to_file") {
+						await writeToFileTool.teardownAbandonedStream(cline)
+					}
+
 					break
 				}
 			}
@@ -797,6 +805,14 @@
 						error.message,
 					)
 
+					// Same leak as the missing-nativeArgs guard: the block streamed partial
+					// deltas (handlePartial ran), this exit bypasses handle()/execute(), so the
+					// per-task stream state and any open diff view must be released here. A
+					// mode file restriction blocking the target path lands exactly here.
+					if (block.name === "write_to_file") {
+						await writeToFileTool.teardownAbandonedStream(cline)
+					}
+
 					break
 				}
 
@@ -873,6 +889,12 @@
 							`Tool call repetition limit reached for ${block.name}. Please try a different approach.`,
 						),
 					)
+
+					// The repetition exit also bypasses handle(): release the stream state and
+					// diff view the partial deltas created, or they leak into the next request.
+					if (block.name === "write_to_file") {
+						await writeToFileTool.teardownAbandonedStream(cline)
+					}
 					break
 				}
 			}