diff --git a/src/__tests__/removeClineFromStack-delegation.spec.ts b/src/__tests__/removeClineFromStack-delegation.spec.ts index bc5b8a4426..b42a549fef 100644 --- a/src/__tests__/removeClineFromStack-delegation.spec.ts +++ b/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 & Partial> & { @@ -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 = { diff --git a/src/core/assistant-message/__tests__/presentAssistantMessage-validation-rejection.spec.ts b/src/core/assistant-message/__tests__/presentAssistantMessage-validation-rejection.spec.ts new file mode 100644 index 0000000000..f656297d97 --- /dev/null +++ b/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[] + userMessageContent: ToolResultBlock[] + didCompleteReadingStream: boolean + didRejectTool: boolean + didAlreadyUseTool: boolean + consecutiveMistakeCount: number + consecutiveMistakeLimit: number + apiConfiguration: { apiProvider: string } + clineMessages: unknown[] + getTaskMode: ReturnType + api: { getModel: () => { id: string; info: Record } } + recordToolUsage: ReturnType + recordToolError: ReturnType + toolRepetitionDetector: { check: ReturnType } + providerRef: { deref: () => { getState: ReturnType } } + say: ReturnType + ask: ReturnType + pushToolResultToUserContent: ReturnType +} + +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() + }) +}) diff --git a/src/core/assistant-message/presentAssistantMessage.ts b/src/core/assistant-message/presentAssistantMessage.ts index 417dcf7a45..baa55963ce 100644 --- a/src/core/assistant-message/presentAssistantMessage.ts +++ b/src/core/assistant-message/presentAssistantMessage.ts @@ -572,6 +572,14 @@ async function presentAssistantMessageBlock(cline: Task): Promise { 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 { 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 { `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 } } diff --git a/src/core/tools/BaseTool.ts b/src/core/tools/BaseTool.ts index 83a733c7b0..339066919e 100644 --- a/src/core/tools/BaseTool.ts +++ b/src/core/tools/BaseTool.ts @@ -98,6 +98,19 @@ export abstract class BaseTool { 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 { + return false + } + /** * Main entry point for tool execution. * @@ -157,7 +170,14 @@ export abstract class BaseTool { } 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 diff --git a/src/core/tools/WriteToFileTool.ts b/src/core/tools/WriteToFileTool.ts index 0c5c80abb9..2cf511cdf2 100644 --- a/src/core/tools/WriteToFileTool.ts +++ b/src/core/tools/WriteToFileTool.ts @@ -1,7 +1,7 @@ import path from "path" import fs from "fs/promises" -import { type ClineSayTool, DEFAULT_WRITE_DELAY_MS } from "@roo-code/types" +import { type ClineSayTool, DEFAULT_WRITE_DELAY_MS, RooCodeEventName } from "@roo-code/types" import { Task } from "../task/Task" import { formatResponse } from "../prompts/responses" @@ -22,18 +22,301 @@ interface WriteToFileParams { content: string } +/** + * Per-task partial-streaming state tracked by WriteToFileTool. + */ +interface TaskPartialStreamState { + /** Last path seen during streaming; undefined until the first delta. */ + lastSeenPartialPath: string | undefined + /** True once a streaming delta hit a fatal filesystem error. */ + streamFailed: boolean + /** The original filesystem error of the failed streaming delta, reported once + * by onParameterParseFailure() when the final block fails to parse (so + * execute() never runs and would never report it). */ + streamError: Error | undefined + /** The task that owns this state; target for abort-listener deregistration. */ + task: Task + /** TaskAborted listener that tears this state down; registered once per task. */ + abortCleanup: () => void +} + export class WriteToFileTool extends BaseTool<"write_to_file"> { readonly name = "write_to_file" as const + /** + * Per-task partial-streaming state, keyed by task id (taskId + instanceId). + * + * All per-task fields live in one object per task so that resetTaskPartialState() / + * resetPartialState() cannot clear a subset of them and leak the rest (abort + * listener, failure mark, path-stabilization entry) for an abandoned stream. + * + * This deliberately diverges from the sibling streaming tools (ApplyDiffTool, + * EditFileTool, SearchReplaceTool, EditTool), which rely on BaseTool's singleton + * lastSeenPartialPath / resetPartialState and keep no failure state. The divergence is + * intentional, for two reasons: + * + * 1. Only this tool's handlePartial performs failure-prone streaming work + * (diffViewProvider.open/update, which can throw EACCES/EROFS); the siblings only + * send a task.ask preview. Without per-task failure tracking, every later delta for + * a failed path would re-attempt the failing operation and re-spawn a partial tool + * message. + * + * 2. The tool instance is a module-level singleton shared by every task, including + * tasks from different ClineProvider instances (e.g. sidebar and tab-panel + * providers, which activate independently). A single provider streams at most one + * task at a time — TaskScheduler gates task.run() at maxConcurrency=1 and + * delegation disposes the parent before the child starts — so per-task keying is + * reachable specifically across providers, where two providers can stream + * write_to_file concurrently through this same singleton. + * + * Lifting this per-task keying into BaseTool for all streaming tools is a follow-up + * (separate PR); it is deliberately not done here. + */ + private taskPartialStreamState = new Map() + + private getPartialStreamFailureKey(task: Task): string { + return `${task.taskId}.${task.instanceId}` + } + + /** + * Get this task's partial stream state, creating it on first use and registering the + * TaskAborted teardown listener exactly once per task. + */ + private getTaskPartialStreamState(task: Task): TaskPartialStreamState { + const key = this.getPartialStreamFailureKey(task) + const existing = this.taskPartialStreamState.get(key) + if (existing) { + return existing + } + + const state: TaskPartialStreamState = { + lastSeenPartialPath: undefined, + streamFailed: false, + streamError: undefined, + task, + abortCleanup: () => this.resetTaskPartialState(task), + } + this.taskPartialStreamState.set(key, state) + task.once(RooCodeEventName.TaskAborted, state.abortCleanup) + return state + } + + private hasPathStabilizedForTask(state: TaskPartialStreamState, partialPath: string | undefined): boolean { + // Stryker disable next-line ConditionalExpression: the `!== undefined` clause is redundant: when + // lastSeenPartialPath is undefined, the second clause only matches an undefined partialPath, which + // the `!!partialPath` in the return value rejects either way -- no test can distinguish the two. + const pathHasStabilized = state.lastSeenPartialPath !== undefined && state.lastSeenPartialPath === partialPath + state.lastSeenPartialPath = partialPath + return pathHasStabilized && !!partialPath + } + + /** + * Clear a task's partial-stream state from a disposal path that does not abort first. + * Task.dispose() removes every listener, so a task disposed directly (for example + * ClineProvider.cleanupFailedHistoryTask()) never fires the TaskAborted cleanup and + * this singleton would keep the disposed task and its diff-view provider. + */ + public clearTaskState(task: Task): void { + this.resetTaskPartialState(task) + } + + /** + * Release what a streamed write left behind when the presenter rejects the COMPLETED + * block before the tool ever runs. + * + * Streaming is not gated by the checks that guard execution: a partial delta registers + * this task's entry (and its TaskAborted listener) and may open a preview, then + * validateToolUse() throws for the completed block - a mode restriction, a disabled + * tool - or the repetition guard refuses it, and the loop breaks before + * writeToFileTool.handle() is reached. None of execute()'s teardown, the parse-failure + * hook, or clearTaskState() runs on that path, so the task would keep the listener, the + * map entry, a preview holding content nobody was asked to approve, and the directories + * that preview created; a retained streamFailed also suppresses this task's later + * previews. + * + * Resources only: the validation error is the tool result the model sees, so nothing is + * reported here beyond a failed rollback, which is a hazard the user - not the model - + * has to know about. A no-op for every task that never streamed. + */ + async releaseStreamAfterValidationRejection(task: Task): Promise { + if (!this.taskPartialStreamState.has(this.getPartialStreamFailureKey(task))) { + return + } + + this.releasePartialStreamBookkeeping(task) + const rollbackError = await this.discardUnapprovedStreamBeforeReset(task) + // Same reason as the rooignore exit below: reset() drops the directories the delta + // adopted without touching disk, so a rejection that never opened a diff view would + // otherwise leave them behind for a write that never happened. + await this.releaseEarlyDirectories(task) + await this.resetDiffViewAfterWrite(task) + if (rollbackError) { + await task + .say( + "error", + "write_to_file: the diff editor could not be restored after the tool call was rejected, so it may still show unapproved content. Do not save that editor.", + ) + .catch((sayError) => { + console.error("Error reporting write_to_file rollback failure:", sayError) + }) + } + } + + /** + * Drop the directories a delta created before the diff view took ownership. + * + * Only for a call that never reached open(): once a session is editing, the provider owns + * those directories and its own discard or revert removes them, so this must not reach + * past the delta that adopted them. + */ + private async releaseEarlyDirectories(task: Task): Promise { + if (task.diffViewProvider.isEditing) { + return + } + await task.diffViewProvider.removeAdoptedDirectories() + } + + private resetTaskPartialState(task: Task): void { + const key = this.getPartialStreamFailureKey(task) + const state = this.taskPartialStreamState.get(key) + if (!state) { + return + } + state.task.off(RooCodeEventName.TaskAborted, state.abortCleanup) + this.taskPartialStreamState.delete(key) + } + + /** + * Release everything the current tool call owns of the partial-stream bookkeeping: + * BaseTool's last-seen path plus THIS task's per-task entry (and its TaskAborted + * listener). Every exit of execute() and every handlePartial() return that skips the + * success-path teardown has to call this. Without it a rejected approval leaves the + * entry attached for the rest of the task's life: the listener is only ever removed + * by a teardown, and a retained streamFailed keeps suppressing this task's later + * diff previews. Task-scoped by design: BaseTool's resetPartialState() only owns the + * singleton last-seen path, and a global teardown would clobber another task that is + * still streaming through this singleton. + */ + private releasePartialStreamBookkeeping(task: Task): void { + super.resetPartialState() + this.resetTaskPartialState(task) + } + + /** + * Whether this task's partial stream is still the live one. handlePartial() awaits + * provider state, a filesystem probe and task.ask() before it touches the diff view; a + * cancellation during any of those awaits runs the TaskAborted teardown (or a direct + * clearTaskState), which deletes this entry. Continuing would re-open a diff view and + * re-ask for a task the user already cancelled, resurrecting the state the teardown + * released. Identity, not presence: a re-created entry for the same key belongs to a new + * stream, and this one must not write into it. + */ + private isPartialStreamStillLive(task: Task, state: TaskPartialStreamState): boolean { + return this.taskPartialStreamState.get(this.getPartialStreamFailureKey(task)) === state + } + + private async resetDiffViewAfterWrite(task: Task): Promise { + await task.diffViewProvider.reset().catch((resetError) => { + console.error("Error resetting write_to_file diff view:", resetError) + }) + } + + /** + * Release a diff view that holds content the user was never asked to approve, and close + * it. + * + * reset() clears the provider's state but leaves the diff document dirty with the streamed + * content; a user save would then persist a write the task never completed (denied or + * failed before approval). Must run BEFORE resetDiffViewAfterWrite(), since reset() clears + * the state the discard relies on. No-op when no diff view is open. + * + * Deliberately NOT revertChanges(): that path SAVES. For a new-file preview it writes the + * partial model output into the placeholder before deleting it, and for a modify it writes + * the restored original back to a file this stream never changed - so a failed save, or a + * failed delete after it, leaves bytes the user never approved on disk. The discard empties + * or restores the buffer in memory only. + * + * A discard failure is RETURNED rather than dropped: the caller records it as the failure + * this stream produced, so debris left on disk is reported instead of being silently + * continued past. + */ + private async discardUnapprovedStreamBeforeReset(task: Task): Promise { + try { + await task.diffViewProvider.discardUnapprovedStream() + } catch (discardError) { + console.error("Error discarding the unapproved write_to_file diff view:", discardError) + return discardError instanceof Error ? discardError : new Error(String(discardError)) + } + return undefined + } + + private async finalizePartialToolAskAfterFailure(task: Task, text?: string): Promise { + await task.finalizePartialToolAsk(text).catch((finalizeError) => { + console.error("Error finalizing write_to_file partial tool ask:", finalizeError) + }) + } + + /** + * Override of BaseTool's teardown boundary for the handle() parse-failure path, where + * execute() never runs and therefore none of execute()'s teardown runs either. + * + * Releases THIS task's stream state: otherwise the abort listener leaks for the task's + * lifetime, and when a streaming delta had failed, the streamFailed guard would suppress + * the diff preview of every later write_to_file in this task. Restores the diff document: + * streaming may have opened it with unapproved partial content, and execute()'s error + * cleanup (revert + reset) never fires on this path, so a user save could persist the + * content without the teardown here. When a streaming delta already hit a fatal filesystem + * error, that error is what the user can act on, so report it with the same "writing file" + * context execute()'s catch uses, and return true to suppress the incidental parse error - + * the failure then surfaces exactly once. + */ + protected override async releaseStreamStateOnParseFailure(task: Task, callbacks: ToolCallbacks): Promise { + const state = this.taskPartialStreamState.get(this.getPartialStreamFailureKey(task)) + if (!state) { + return false + } + + this.resetTaskPartialState(task) + const rollbackError = await this.discardUnapprovedStreamBeforeReset(task) + await this.resetDiffViewAfterWrite(task) + + // A failed rollback is the more actionable failure (debris is still on disk), so it takes + // the report slot when present; the streaming error is kept behind it as the cause. + if (rollbackError) { + await callbacks.handleError( + "writing file", + new Error(`write_to_file rollback failed after a streaming error: ${rollbackError.message}`, { + cause: state.streamError ?? rollbackError, + }), + ) + return true + } + + if (state.streamError) { + await callbacks.handleError("writing file", state.streamError) + return true + } + + return false + } + async execute(params: WriteToFileParams, task: Task, callbacks: ToolCallbacks): Promise { const { pushToolResult, handleError, askApproval } = callbacks const relPath = params.path let newContent = params.content + // Set when this execute() opens its own partial tool ask (diff-view branch), so + // the catch below can finalize it. Undefined on the saveDirectly branch. + let pendingPartialAsk: string | undefined if (!relPath) { task.consecutiveMistakeCount++ task.recordToolError("write_to_file") pushToolResult(await task.sayAndCreateMissingParamError("write_to_file", "path")) + + // Returning here skips the try/catch teardown below: release THIS task's stream + // state (and only this task's) so the abort listener and any streamFailed guard do + // not outlive the call. + this.releasePartialStreamBookkeeping(task) await task.diffViewProvider.reset() return } @@ -42,6 +325,11 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { task.consecutiveMistakeCount++ task.recordToolError("write_to_file") pushToolResult(await task.sayAndCreateMissingParamError("write_to_file", "content")) + + // Returning here skips the try/catch teardown below: release THIS task's stream + // state (and only this task's) so the abort listener and any streamFailed guard do + // not outlive the call. + this.releasePartialStreamBookkeeping(task) await task.diffViewProvider.reset() return } @@ -51,48 +339,92 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { if (!accessAllowed) { await task.say("rooignore_error", relPath) pushToolResult(formatResponse.rooIgnoreError(relPath)) + + // Returning here skips the try/catch teardown below: release THIS task's stream + // state (and only this task's) so the abort listener and any streamFailed guard do + // not outlive the call. + this.releasePartialStreamBookkeeping(task) + // Streaming is not gated by this check - open() never consults rooignore - so the + // denied call can still be holding a preview full of content that will never be + // approved. Leaving it open hands the next write a live editor containing someone + // else's unapproved content. + if (task.diffViewProvider.isEditing) { + const discardError = await this.discardUnapprovedStreamBeforeReset(task) + if (discardError) { + await task + .say( + "error", + `write_to_file could not discard the preview for the denied write: ${discardError.message}`, + ) + .catch((sayError) => { + console.error("Error reporting write_to_file discard failure:", sayError) + }) + } + } + // A delta that adopted parent directories before any diff view existed still owns + // them here: reset() drops the recorded list without touching disk, so on the no-editor + // path the directories of a denied write would stay on disk. When a session IS editing + // the provider owns them and the discard above has already removed them, which is why + // this runs after it and guards on isEditing itself. + await this.releaseEarlyDirectories(task) + await this.resetDiffViewAfterWrite(task) return } const isWriteProtected = task.rooProtectedController?.isWriteProtected(relPath) || false + // Declared outside the setup boundary below: the write body's try block reads it after the + // setup succeeded, and a block-scoped declaration would not be visible there. let fileExists: boolean - const absolutePath = path.resolve(task.cwd, relPath) + let sharedMessageProps: ClineSayTool + + // The setup below awaits before the write's own try/catch begins. A rejection there (an + // EACCES/EROFS from createDirectoriesForFile, a failing filesystem probe) would otherwise + // escape execute() with this task's stream entry and its TaskAborted listener still + // registered, so the map would keep the task alive and a later write could reuse a stale + // stream state. Release and rethrow: BaseTool.handle() still reports the error exactly + // once, and no diff view is open yet, so there is no preview to discard. + try { + const absolutePath = path.resolve(task.cwd, relPath) - if (task.diffViewProvider.editType !== undefined) { - fileExists = task.diffViewProvider.editType === "modify" - } else { - fileExists = await fileExistsAtPath(absolutePath) - task.diffViewProvider.editType = fileExists ? "modify" : "create" - } + if (task.diffViewProvider.editType !== undefined) { + fileExists = task.diffViewProvider.editType === "modify" + } else { + fileExists = await fileExistsAtPath(absolutePath) + task.diffViewProvider.editType = fileExists ? "modify" : "create" + } - // Create parent directories early for new files to prevent ENOENT errors - // in subsequent operations (e.g., diffViewProvider.open, fs.readFile) - if (!fileExists) { - await createDirectoriesForFile(absolutePath) - } + // Create parent directories early for new files to prevent ENOENT errors + // in subsequent operations (e.g., diffViewProvider.open, fs.readFile) + if (!fileExists) { + await createDirectoriesForFile(absolutePath) + } - if (newContent.startsWith("```")) { - newContent = newContent.split("\n").slice(1).join("\n") - } + if (newContent.startsWith("```")) { + newContent = newContent.split("\n").slice(1).join("\n") + } - if (newContent.endsWith("```")) { - newContent = newContent.split("\n").slice(0, -1).join("\n") - } + if (newContent.endsWith("```")) { + newContent = newContent.split("\n").slice(0, -1).join("\n") + } - if (!task.api.getModel().id.includes("claude")) { - newContent = unescapeHtmlEntities(newContent) - } + if (!task.api.getModel().id.includes("claude")) { + newContent = unescapeHtmlEntities(newContent) + } - const fullPath = relPath ? path.resolve(task.cwd, relPath) : "" - const isOutsideWorkspace = isPathOutsideWorkspace(fullPath) + const fullPath = relPath ? path.resolve(task.cwd, relPath) : "" + const isOutsideWorkspace = isPathOutsideWorkspace(fullPath) - const sharedMessageProps: ClineSayTool = { - tool: fileExists ? "editedExistingFile" : "newFileCreated", - path: getReadablePath(task.cwd, relPath), - content: newContent, - isOutsideWorkspace, - isProtected: isWriteProtected, + sharedMessageProps = { + tool: fileExists ? "editedExistingFile" : "newFileCreated", + path: getReadablePath(task.cwd, relPath), + content: newContent, + isOutsideWorkspace, + isProtected: isWriteProtected, + } + } catch (error) { + this.releasePartialStreamBookkeeping(task) + throw error } try { @@ -129,6 +461,7 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { const didApprove = await askApproval("tool", completeMessage, undefined, isWriteProtected) if (!didApprove) { + this.releasePartialStreamBookkeeping(task) return } @@ -136,6 +469,7 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { } else { if (!task.diffViewProvider.isEditing) { const partialMessage = JSON.stringify(sharedMessageProps) + pendingPartialAsk = partialMessage await task.ask("tool", partialMessage, true).catch(() => {}) await task.diffViewProvider.open(relPath) } @@ -161,6 +495,7 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { if (!didApprove) { await task.diffViewProvider.revertChanges() + this.releasePartialStreamBookkeeping(task) return } @@ -178,15 +513,47 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { pushToolResult(message) await task.diffViewProvider.reset() - this.resetPartialState() + // Tear down only this task's entry; clearing the whole map would drop another + // task's streamFailed/streamError while it is still streaming. + this.releasePartialStreamBookkeeping(task) task.processQueuedMessages() return } catch (error) { + // The diff-view branch above may have opened a fresh partial ask for this + // (retried) write. Finalize it before tearing down, or the spinner and + // Save/Reject buttons stay live for a tool call that has already failed. + if (pendingPartialAsk !== undefined) { + await this.finalizePartialToolAskAfterFailure(task, pendingPartialAsk) + } await handleError("writing file", error as Error) - await task.diffViewProvider.reset() - this.resetPartialState() + // A preview that never got its approved write to disk holds unapproved content: discard + // it before reset() drops the state the discard depends on. This is also the teardown for + // a saveChanges() that rejected mid-write - the placeholder is still owned by this edit, + // because saveChanges() releases ownership only once the save lands - so the discard is + // what removes the empty or half-written new file. + if (task.diffViewProvider.isEditing) { + const discardError = await this.discardUnapprovedStreamBeforeReset(task) + if (discardError) { + // The report must not own the teardown: Task.say() throws once the task is + // aborted, and a rejection here would skip the reset() and the bookkeeping + // release below - leaking exactly what this block exists to clean up. + await task + .say( + "error", + `write_to_file could not discard the unapproved preview after the failed write: ${discardError.message}`, + ) + .catch((sayError) => { + console.error("Error reporting write_to_file discard failure:", sayError) + }) + } + } + // Guarded: a reset() that rejects must not skip the release below, and it must + // not turn a reported write failure into a teardown failure the caller cannot act + // on - the per-task entry and its TaskAborted listener are still ours to drop. + await this.resetDiffViewAfterWrite(task) + this.releasePartialStreamBookkeeping(task) return } } @@ -195,62 +562,180 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { const relPath: string | undefined = block.params.path const newContent: string | undefined = block.params.content - // Wait for path to stabilize before showing UI (prevents truncated paths) - if (!this.hasPathStabilized(relPath) || newContent === undefined) { + // Get (or create) this task's state; registers the TaskAborted teardown listener + // once, so abandoned streams are torn down even if execute() never runs. + const partialStreamState = this.getTaskPartialStreamState(task) + + // A delta of THIS call already failed at the diff view. Retrying on every later delta + // would re-open a diff editor that just failed and re-spawn a partial tool message for a + // call that has already reported its error, so the rest of the stream is suppressed until + // whichever teardown ends the call releases the entry. + if (partialStreamState.streamFailed) { return } - const provider = task.providerRef.deref() - const state = await provider?.getState() - const isPreventFocusDisruptionEnabled = experiments.isEnabled( - state?.experiments ?? {}, - EXPERIMENT_IDS.PREVENT_FOCUS_DISRUPTION, - ) - - if (isPreventFocusDisruptionEnabled) { + // Wait for path to stabilize before showing UI (prevents truncated paths) + if (!this.hasPathStabilizedForTask(partialStreamState, relPath) || newContent === undefined) { return } - // relPath is guaranteed non-null after hasPathStabilized - let fileExists: boolean - const absolutePath = path.resolve(task.cwd, relPath!) + // Which failure window this delta is in, read by the catch below: false while the + // pre-streaming setup runs, true from the moment the diff view is the thing at risk. Declared + // outside the try because a try block's own scope is not visible to its catch. + let diffViewStarted = false - if (task.diffViewProvider.editType !== undefined) { - fileExists = task.diffViewProvider.editType === "modify" - } else { - fileExists = await fileExistsAtPath(absolutePath) - task.diffViewProvider.editType = fileExists ? "modify" : "create" - } + try { + // Everything from here up to the diff view is setup that can fail before + // execute() ever runs; the catch below owns the teardown for that window. + const provider = task.providerRef.deref() + const state = await provider?.getState() - // Create parent directories early for new files to prevent ENOENT errors - // in subsequent operations (e.g., diffViewProvider.open) - if (!fileExists) { - await createDirectoriesForFile(absolutePath) - } + // Cancelled while provider state was in flight: the teardown already + // released this task's stream state. + if (!this.isPartialStreamStillLive(task, partialStreamState)) { + return + } + const isPreventFocusDisruptionEnabled = experiments.isEnabled( + state?.experiments ?? {}, + EXPERIMENT_IDS.PREVENT_FOCUS_DISRUPTION, + ) + + if (isPreventFocusDisruptionEnabled) { + // The preview is suppressed for this stream: release the entry registered above so + // the abort listener and any failure mark do not outlive a delta that never shows + // a diff view and never reaches execute()'s teardown. + this.releasePartialStreamBookkeeping(task) + return + } - const isWriteProtected = task.rooProtectedController?.isWriteProtected(relPath!) || false - const isOutsideWorkspace = isPathOutsideWorkspace(absolutePath) + // relPath is guaranteed non-null after hasPathStabilized + let fileExists: boolean + const absolutePath = path.resolve(task.cwd, relPath!) - const sharedMessageProps: ClineSayTool = { - tool: fileExists ? "editedExistingFile" : "newFileCreated", - path: getReadablePath(task.cwd, relPath!), - content: newContent || "", - isOutsideWorkspace, - isProtected: isWriteProtected, - } + if (task.diffViewProvider.editType !== undefined) { + fileExists = task.diffViewProvider.editType === "modify" + } else { + fileExists = await fileExistsAtPath(absolutePath) + if (!this.isPartialStreamStillLive(task, partialStreamState)) { + return + } + task.diffViewProvider.editType = fileExists ? "modify" : "create" + } - const partialMessage = JSON.stringify(sharedMessageProps) - await task.ask("tool", partialMessage, block.partial).catch(() => {}) + // Create parent directories early for new files to prevent ENOENT errors + // in subsequent operations (e.g., diffViewProvider.open) + if (!fileExists) { + // Handed to the diff view's cleanup state at once, not at open(): open() records + // only the directories it creates itself, which is none after this call, and every + // teardown removes what that list holds. + task.diffViewProvider.adoptCreatedDirectories(await createDirectoriesForFile(absolutePath)) + } + // Abandonment can land while the directory creation is in flight: its teardown has + // already released this task's stream state, and asking or streaming now would show a + // partial tool call for a task that no longer exists. + if (!this.isPartialStreamStillLive(task, partialStreamState)) { + await this.releaseEarlyDirectories(task) + return + } + + const isWriteProtected = task.rooProtectedController?.isWriteProtected(relPath!) || false + const isOutsideWorkspace = isPathOutsideWorkspace(absolutePath) - if (newContent) { - if (!task.diffViewProvider.isEditing) { - await task.diffViewProvider.open(relPath!) + const sharedMessageProps: ClineSayTool = { + tool: fileExists ? "editedExistingFile" : "newFileCreated", + path: getReadablePath(task.cwd, relPath!), + content: newContent || "", + isOutsideWorkspace, + isProtected: isWriteProtected, } - await task.diffViewProvider.update( - everyLineHasLineNumbers(newContent) ? stripLineNumbers(newContent) : newContent, - false, - ) + const partialMessage = JSON.stringify(sharedMessageProps) + await task.ask("tool", partialMessage, block.partial).catch(() => {}) + + if (!this.isPartialStreamStillLive(task, partialStreamState)) { + await this.releaseEarlyDirectories(task) + return + } + + // From here the diff view is what can fail, and its failure has to be owned by this + // boundary rather than escaping to BaseTool's generic catch, which reports the error but + // releases nothing: the entry and its TaskAborted listener would stay attached and the + // next delta would retry the operation that just failed. + diffViewStarted = true + + if (newContent) { + if (!task.diffViewProvider.isEditing) { + await task.diffViewProvider.open(relPath!) + } + + // Cancellation may land while open() is in flight: its abort handler has + // already torn the stream down (and may have reverted or closed this very + // diff view), so streaming the partial content into it now would resurrect a + // view for a task that no longer exists. + if (!this.isPartialStreamStillLive(task, partialStreamState)) { + // open() can finish after the abort cleanup already ran, and that cleanup checked + // isEditing at a moment when there was no session yet. The view and the new-file + // placeholder open() wrote are then open with nobody left to close them, so this + // delta - the one that opened them - owns their teardown. + if (task.diffViewProvider.isEditing) { + await this.discardUnapprovedStreamBeforeReset(task) + await this.resetDiffViewAfterWrite(task) + } + await this.releaseEarlyDirectories(task) + return + } + + await task.diffViewProvider.update( + everyLineHasLineNumbers(newContent) ? stripLineNumbers(newContent) : newContent, + false, + ) + } + } catch (error) { + // Two windows, two owners of the teardown. + // + // * Pre-streaming setup (provider state, the filesystem probe, directory creation, the + // partial ask): nothing was shown and nothing of ours is on disk, so this call's entry + // is released - a later delta may legitimately retry a transient setup failure. + // * The diff view (open()/update()): the preview itself is broken for this call. Releasing + // here would let the next delta re-register and re-open the view that just failed, which + // is the retry this boundary exists to stop, so the entry is kept and marked failed + // instead: the guard at the top of handlePartial suppresses the rest of the stream, and + // whichever teardown ends the call - the parse-failure boundary, which reports the + // recorded error, execute(), or a cancellation - releases it and its listener. + // + // Either way the unapproved preview is discarded and the view reset before the error is + // rethrown, so BaseTool.handle() still reports it exactly once. + if (diffViewStarted) { + partialStreamState.streamFailed = true + partialStreamState.streamError = error instanceof Error ? error : new Error(String(error)) + } else { + this.releasePartialStreamBookkeeping(task) + // The setup window owns whatever the diff view never took over, including the + // directories this delta created a few lines above. + await this.releaseEarlyDirectories(task) + } + if (task.diffViewProvider.isEditing) { + const discardError = await this.discardUnapprovedStreamBeforeReset(task) + await this.resetDiffViewAfterWrite(task) + if (discardError) { + // Two failures, two channels. The exception this delta produced stays the primary one: + // BaseTool.handle() reports it unchanged, so wrapping it here would change the failure + // the caller sees. The discard failure is a separate, actionable condition - the + // placeholder or the created directories are still on disk and the buffer may still hold + // unapproved content - so it gets its own report instead of only a console line. + // A rejected report must not swallow the exception this delta produced: the + // throw below is the caller's contract. + await task + .say( + "error", + `write_to_file could not discard the unapproved preview after the failed stream: ${discardError.message}`, + ) + .catch((sayError) => { + console.error("Error reporting write_to_file discard failure:", sayError) + }) + } + } + throw error } } } diff --git a/src/core/tools/__tests__/BaseTool-parse-failure-default-hook.spec.ts b/src/core/tools/__tests__/BaseTool-parse-failure-default-hook.spec.ts new file mode 100644 index 0000000000..028122c258 --- /dev/null +++ b/src/core/tools/__tests__/BaseTool-parse-failure-default-hook.spec.ts @@ -0,0 +1,59 @@ +// npx vitest run core/tools/__tests__/BaseTool-parse-failure-default-hook.spec.ts + +import { describe, it, expect, vi } from "vitest" + +import { BaseTool, type ToolCallbacks } from "../BaseTool" +import type { Task } from "../../task/Task" +import type { ToolUse } from "../../../shared/tools" + +/** + * A tool that keeps no per-task stream state: it inherits BaseTool's default + * releaseStreamStateOnParseFailure(), which reports nothing and returns false. Every tool + * except write_to_file runs through that default, so the generic parse error is the only + * thing a user would ever see for a malformed call - if the default ever started + * swallowing it, parse errors would vanish for the whole tool set at once. + */ +class DefaultHookTool extends BaseTool<"execute_command"> { + readonly name = "execute_command" as const + execute = vi.fn().mockResolvedValue(undefined) +} + +const buildTask = (): Task => ({ taskId: "parse-failure-task", instanceId: "inst-1" }) as unknown as Task + +const buildCallbacks = (): ToolCallbacks => ({ + askApproval: vi.fn().mockResolvedValue({ response: "yesButtonClicked" }), + handleError: vi.fn().mockResolvedValue(undefined), + pushToolResult: vi.fn(), +}) + +describe("BaseTool parse-failure path with the default stream-release hook", () => { + it("skips execute() and reports the generic parse error exactly once", async () => { + const tool = new DefaultHookTool() + const task = buildTask() + const callbacks = buildCallbacks() + // A non-native block: no nativeArgs, and params carry no XML markup, so the parse step + // fails on the missing native arguments rather than on the legacy-format message. + const block = { + type: "tool_use", + id: "call-parse-1", + name: "execute_command", + params: { command: "ls" }, + partial: false, + } as unknown as ToolUse<"execute_command"> + + await tool.handle(task, block, callbacks) + + expect(tool.execute).not.toHaveBeenCalled() + expect(callbacks.handleError).toHaveBeenCalledTimes(1) + expect(callbacks.handleError).toHaveBeenCalledWith( + "parsing execute_command args", + expect.objectContaining({ message: expect.stringContaining("missing native arguments") }), + ) + // The hook consulted here is BaseTool's own, not an override: this is the path every + // tool other than write_to_file takes. + expect(DefaultHookTool.prototype["releaseStreamStateOnParseFailure"]).toBe( + BaseTool.prototype["releaseStreamStateOnParseFailure"], + ) + expect(callbacks.pushToolResult).not.toHaveBeenCalled() + }) +}) diff --git a/src/core/tools/__tests__/writeToFileTool-partial-state-cleanup.spec.ts b/src/core/tools/__tests__/writeToFileTool-partial-state-cleanup.spec.ts new file mode 100644 index 0000000000..afd9987fd9 --- /dev/null +++ b/src/core/tools/__tests__/writeToFileTool-partial-state-cleanup.spec.ts @@ -0,0 +1,186 @@ +// npx vitest run core/tools/__tests__/writeToFileTool-partial-state-cleanup.spec.ts + +import { RooCodeEventName } from "@roo-code/types" +import { vi, type MockedFunction } from "vitest" + +import { type Task } from "../../task/Task" +import { writeToFileTool } from "../WriteToFileTool" + +// The cleanup primitives only read these members, so a structural double is enough; +// the double assertion is the repo's existing pattern for private-method tests +// (see src/__tests__/removeClineFromStack-delegation.spec.ts). +interface CleanupTask { + taskId: string + instanceId: string + once: MockedFunction<(...args: unknown[]) => unknown> + off: MockedFunction<(...args: unknown[]) => unknown> + diffViewProvider: { + reset: MockedFunction<() => Promise> + revertChanges: MockedFunction<() => Promise> + discardUnapprovedStream: MockedFunction<() => Promise> + adoptCreatedDirectories: MockedFunction<(directories: string[]) => void> + removeAdoptedDirectories: MockedFunction<() => Promise> + } + finalizePartialToolAsk: MockedFunction<() => Promise> + say: MockedFunction<() => Promise> +} + +function buildTask(taskId: string, instanceId: string): Task { + const task: CleanupTask = { + taskId, + instanceId, + once: vi.fn(), + off: vi.fn(), + diffViewProvider: { + reset: vi.fn().mockResolvedValue(undefined), + revertChanges: vi.fn().mockResolvedValue(undefined), + discardUnapprovedStream: vi.fn().mockResolvedValue(undefined), + adoptCreatedDirectories: vi.fn(), + removeAdoptedDirectories: vi.fn().mockResolvedValue(undefined), + }, + finalizePartialToolAsk: vi.fn().mockResolvedValue(undefined), + say: vi.fn().mockResolvedValue(undefined), + } + return task as unknown as Task +} + +// Private members are reached by bracket notation (AGENTS.md: no `as any`). +const stateFor = (task: Task) => writeToFileTool["taskPartialStreamState"].get(`${task.taskId}.${task.instanceId}`) + +describe("WriteToFileTool per-task partial-state cleanup", () => { + afterEach(() => { + writeToFileTool["taskPartialStreamState"].clear() + vi.restoreAllMocks() + }) + + it("releases the task state and deregisters the abort listener", async () => { + const task = buildTask("cleanup-task", "inst-1") + const state = writeToFileTool["getTaskPartialStreamState"](task) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + + writeToFileTool.clearTaskState(task) + + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + expect((task as unknown as CleanupTask).off).toHaveBeenCalledWith( + RooCodeEventName.TaskAborted, + state.abortCleanup, + ) + }) + + it("is a no-op for a task that never streamed", async () => { + const task = buildTask("never-streamed", "inst-2") + + writeToFileTool.clearTaskState(task) + + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + expect((task as unknown as CleanupTask).off).not.toHaveBeenCalled() + }) + + it("logs and continues when resetting the diff view fails", async () => { + const task = buildTask("reset-fails", "inst-3") + const t = task as unknown as CleanupTask + t.diffViewProvider.reset = vi.fn().mockRejectedValue(new Error("reset failed")) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + await writeToFileTool["resetDiffViewAfterWrite"](task) + + expect(errorSpy).toHaveBeenCalledWith("Error resetting write_to_file diff view:", expect.any(Error)) + }) + + it("logs and continues when discarding the unapproved diff view fails", async () => { + const task = buildTask("discard-fails", "inst-4") + const t = task as unknown as CleanupTask + t.diffViewProvider.discardUnapprovedStream = vi.fn().mockRejectedValue(new Error("discard failed")) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + const failure = await writeToFileTool["discardUnapprovedStreamBeforeReset"](task) + + expect(errorSpy).toHaveBeenCalledWith( + "Error discarding the unapproved write_to_file diff view:", + expect.any(Error), + ) + // Returned, not dropped: the caller reports the debris instead of continuing past it. + expect(failure?.message).toBe("discard failed") + }) + + it("releases the stream state and the preview when the completed block is rejected before the tool runs", async () => { + // Streaming is not gated by the checks that guard execution: a partial delta registers + // this entry and may open a preview, then validateToolUse() throws (or the repetition + // guard refuses the block) and the loop breaks before handle() is reached. Nothing else + // on that path releases the entry, the listener, or the preview. + const task = buildTask("validation-rejected", "inst-6") + const t = task as unknown as CleanupTask + const state = writeToFileTool["getTaskPartialStreamState"](task) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + + await writeToFileTool.releaseStreamAfterValidationRejection(task) + + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + expect(t.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, state.abortCleanup) + expect(t.diffViewProvider.discardUnapprovedStream).toHaveBeenCalledTimes(1) + expect(t.diffViewProvider.reset).toHaveBeenCalledTimes(1) + // The discard must run first: reset() clears the state it reads. + expect(t.diffViewProvider.discardUnapprovedStream.mock.invocationCallOrder[0]).toBeLessThan( + t.diffViewProvider.reset.mock.invocationCallOrder[0], + ) + // Resources only - the validation error is the tool result the model sees. + expect(t.say).not.toHaveBeenCalled() + }) + + it("is a no-op when the rejected block never streamed", async () => { + // Every other tool and every non-streaming write reaches this branch. + const task = buildTask("rejected-never-streamed", "inst-7") + const t = task as unknown as CleanupTask + + await writeToFileTool.releaseStreamAfterValidationRejection(task) + + expect(t.off).not.toHaveBeenCalled() + expect(t.diffViewProvider.discardUnapprovedStream).not.toHaveBeenCalled() + expect(t.diffViewProvider.reset).not.toHaveBeenCalled() + }) + + it("tells the user the editor may still hold unapproved content when that rejection cannot restore it", async () => { + // The validation error is the model's tool result and must stay the only one; the disk + // hazard belongs to the user, so it surfaces in the chat instead of replacing it. + const task = buildTask("rejected-rollback-fails", "inst-8") + const t = task as unknown as CleanupTask + t.diffViewProvider.discardUnapprovedStream = vi.fn().mockRejectedValue(new Error("close rejected")) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + writeToFileTool["getTaskPartialStreamState"](task) + + await writeToFileTool.releaseStreamAfterValidationRejection(task) + + expect(t.say).toHaveBeenCalledWith("error", expect.stringContaining("may still show unapproved content")) + expect(t.diffViewProvider.reset).toHaveBeenCalledTimes(1) + errorSpy.mockRestore() + }) + + it("resolves and logs when the rollback report itself cannot be delivered", async () => { + // Task.say() throws once the task is aborted, and the presenter can reject a block + // while that abort lands. The release must still finish and reset the view: a report + // that cannot be delivered is not a reason to stop cleaning up. + const task = buildTask("rejected-report-fails", "inst-9") + const t = task as unknown as CleanupTask + t.diffViewProvider.discardUnapprovedStream = vi.fn().mockRejectedValue(new Error("close rejected")) + t.say = vi.fn().mockRejectedValue(new Error("task aborted")) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + writeToFileTool["getTaskPartialStreamState"](task) + + await expect(writeToFileTool.releaseStreamAfterValidationRejection(task)).resolves.toBeUndefined() + + expect(errorSpy).toHaveBeenCalledWith("Error reporting write_to_file rollback failure:", expect.any(Error)) + expect(t.diffViewProvider.reset).toHaveBeenCalledTimes(1) + errorSpy.mockRestore() + }) + + it("logs and continues when finalizing the open partial ask fails", async () => { + const task = buildTask("finalize-fails", "inst-5") + const t = task as unknown as CleanupTask + t.finalizePartialToolAsk = vi.fn().mockRejectedValue(new Error("finalize failed")) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + await writeToFileTool["finalizePartialToolAskAfterFailure"](task, "partial text") + + expect(errorSpy).toHaveBeenCalledWith("Error finalizing write_to_file partial tool ask:", expect.any(Error)) + }) +}) diff --git a/src/core/tools/__tests__/writeToFileTool.spec.ts b/src/core/tools/__tests__/writeToFileTool.spec.ts index 52a7e3c052..dd82c1b001 100644 --- a/src/core/tools/__tests__/writeToFileTool.spec.ts +++ b/src/core/tools/__tests__/writeToFileTool.spec.ts @@ -1,5 +1,6 @@ import * as path from "path" +import { RooCodeEventName } from "@roo-code/types" import type { MockedFunction } from "vitest" import { fileExistsAtPath, createDirectoriesForFile } from "../../../utils/fs" @@ -115,9 +116,18 @@ describe("writeToFileTool", () => { beforeEach(() => { vi.clearAllMocks() writeToFileTool.resetPartialState() + // Per-task entries are released by the tool's own teardown paths (execute() exits, the + // handle() parse-failure hook, clearTaskState). The suite clears them explicitly so no + // test inherits another test's stream state or abort listener. + for (const state of [...writeToFileTool["taskPartialStreamState"].values()]) { + writeToFileTool.clearTaskState(state.task) + } mockedPathResolve.mockReturnValue(absoluteFilePath) mockedFileExistsAtPath.mockResolvedValue(false) + // vi.clearAllMocks() keeps the last mock implementation; reset the factory default here + // so no test depends on declaration order or an earlier test's rejection. + mockedCreateDirectoriesForFile.mockResolvedValue([]) mockedIsPathOutsideWorkspace.mockReturnValue(false) mockedGetReadablePath.mockReturnValue("test/path.txt") mockedUnescapeHtmlEntities.mockImplementation((content) => { @@ -128,6 +138,8 @@ describe("writeToFileTool", () => { return content }) + mockCline.taskId = "task-1" + mockCline.instanceId = "instance-1" mockCline.cwd = "/" mockCline.consecutiveMistakeCount = 0 mockCline.didEditFile = false @@ -151,6 +163,9 @@ describe("writeToFileTool", () => { update: vi.fn().mockResolvedValue(undefined), reset: vi.fn().mockResolvedValue(undefined), revertChanges: vi.fn().mockResolvedValue(undefined), + discardUnapprovedStream: vi.fn().mockResolvedValue(undefined), + adoptCreatedDirectories: vi.fn(), + removeAdoptedDirectories: vi.fn().mockResolvedValue(undefined), saveChanges: vi.fn().mockResolvedValue({ newProblemsMessage: "", userEdits: null, @@ -186,8 +201,12 @@ describe("writeToFileTool", () => { } mockCline.say = vi.fn().mockResolvedValue(undefined) mockCline.ask = vi.fn().mockResolvedValue(undefined) + mockCline.once = vi.fn() + mockCline.off = vi.fn() + mockCline.finalizePartialToolAsk = vi.fn().mockResolvedValue(undefined) mockCline.recordToolError = vi.fn() mockCline.sayAndCreateMissingParamError = vi.fn().mockResolvedValue("Missing param error") + mockCline.processQueuedMessages = vi.fn() mockAskApproval = vi.fn().mockResolvedValue(true) mockHandleError = vi.fn().mockResolvedValue(undefined) @@ -419,6 +438,1088 @@ describe("writeToFileTool", () => { expect(mockCline.diffViewProvider.open).toHaveBeenCalledWith(testFilePath) expect(mockCline.diffViewProvider.update).toHaveBeenCalledWith(testContent, false) }) + + it("cleans per-task partial state when the task aborts before execute finalization", async () => { + let abortCleanup: (() => void) | undefined + mockCline.once.mockImplementation((event: RooCodeEventName, listener: () => void) => { + if (event === RooCodeEventName.TaskAborted) { + abortCleanup = listener + } + return mockCline + }) + + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + expect(mockCline.ask).toHaveBeenCalledTimes(1) + // One listener per task, not one per delta, and the teardown must deregister THAT function: + // expect.any(Function) would also pass a tool that registered twice or removed a + // different callback and left the real listener attached. + const registrations = mockCline.once.mock.calls.filter( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + ) + expect(registrations).toHaveLength(1) + expect(mockCline.once).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortCleanup) + + abortCleanup?.() + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortCleanup) + + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + expect(mockCline.ask).toHaveBeenCalledTimes(1) + + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + expect(mockCline.ask).toHaveBeenCalledTimes(2) + }) + + it("does not treat a changed path between deltas as stabilized", async () => { + // Delta 1 streams "alpha.txt"; delta 2 streams "beta.txt" for the same task. The path changed + // between deltas, so it must not count as stabilized and no partial `tool` ask may be issued for + // the still-changing second path. + await executeWriteFileTool({ path: "alpha.txt" }, { isPartial: true }) + await executeWriteFileTool({ path: "beta.txt" }, { isPartial: true }) + + expect(mockCline.ask).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.open).not.toHaveBeenCalled() + }) + }) + + describe("path stabilization predicate", () => { + // The predicate is exercised directly (it is private) because not all of its branches are + // observable through handlePartial(): an undefined path reaches the same early return either + // way, so the clause-by-clause behavior must be pinned at the predicate level. + function makeState(lastSeenPartialPath: string | undefined) { + return { + lastSeenPartialPath, + streamFailed: false, + streamError: undefined, + task: mockCline, + abortCleanup: () => {}, + } + } + + it("reports a first delta as not stabilized and records the seen path", () => { + const state = makeState(undefined) + + expect(writeToFileTool["hasPathStabilizedForTask"](state, "a.txt")).toBe(false) + expect(state.lastSeenPartialPath).toBe("a.txt") + }) + + it("reports a repeated path as stabilized", () => { + const state = makeState("a.txt") + + expect(writeToFileTool["hasPathStabilizedForTask"](state, "a.txt")).toBe(true) + }) + + it("reports a changed path as not stabilized", () => { + const state = makeState("a.txt") + + expect(writeToFileTool["hasPathStabilizedForTask"](state, "b.txt")).toBe(false) + expect(state.lastSeenPartialPath).toBe("b.txt") + }) + }) + + describe("resetPartialState", () => { + it("resets only the singleton path and leaves every task's stream state alone", async () => { + // The tool instance is a module-level singleton shared by concurrent tasks, so the + // base-class reset owns only lastSeenPartialPath. Clearing every task's entry from here + // would drop another task's streamFailed/streamError while it is still streaming; the + // per-task teardown (execute() exits, the parse-failure hook, clearTaskState) owns those. + let abortCleanup: (() => void) | undefined + mockCline.once.mockImplementation((event: RooCodeEventName, listener: () => void) => { + if (event === RooCodeEventName.TaskAborted) { + abortCleanup = listener + } + return mockCline + }) + + // Seed one per-task state with an abort listener attached. + await executeWriteFileTool({}, { isPartial: true }) + await executeWriteFileTool({}, { isPartial: true }) + expect(mockCline.ask).toHaveBeenCalledTimes(1) + expect(abortCleanup).toBeTypeOf("function") + + writeToFileTool["lastSeenPartialPath"] = "stale-path" + writeToFileTool.resetPartialState() + + expect(writeToFileTool["lastSeenPartialPath"]).toBeUndefined() + expect(mockCline.off).not.toHaveBeenCalled() + // The entry survives, and with it the per-task path stabilization: the next delta is + // still the same live stream, so it goes straight to the partial ask instead of + // restarting an un-stabilized sequence. + await executeWriteFileTool({}, { isPartial: true }) + expect(mockCline.ask).toHaveBeenCalledTimes(2) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + + // The task-scoped teardown is what releases it. + writeToFileTool.clearTaskState(mockCline) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortCleanup) + }) + }) + + describe("per-task stream state isolation", () => { + // A second task streaming through the same singleton while mockCline's execute() + // runs. Structural double, same pattern as the partial-state-cleanup spec. + function buildStreamingTask(taskId: string, instanceId: string) { + return { + taskId, + instanceId, + once: vi.fn(), + off: vi.fn(), + diffViewProvider: { + reset: vi.fn().mockResolvedValue(undefined), + revertChanges: vi.fn().mockResolvedValue(undefined), + discardUnapprovedStream: vi.fn().mockResolvedValue(undefined), + adoptCreatedDirectories: vi.fn(), + removeAdoptedDirectories: vi.fn().mockResolvedValue(undefined), + }, + finalizePartialToolAsk: vi.fn().mockResolvedValue(undefined), + } + } + + it("leaves another task's stream state intact when execute() completes", async () => { + const other = buildStreamingTask("task-2", "instance-2") + const otherState = writeToFileTool["getTaskPartialStreamState"](other as never) + otherState.streamFailed = true + otherState.streamError = new Error("other task stream failure") + + await executeWriteFileTool({}) + + // The other task is still streaming: its failure state must survive, or its + // next delta re-opens the diff view and spawns a duplicate partial ask. + const retained = writeToFileTool["taskPartialStreamState"].get("task-2.instance-2") + expect(retained).toBeDefined() + expect(retained?.streamFailed).toBe(true) + expect(retained?.streamError?.message).toBe("other task stream failure") + expect(other.off).not.toHaveBeenCalled() + }) + + it("finalizes the partial ask when the write itself fails", async () => { + // Exact payload streamed as the partial tool ask for this scenario; a weaker + // matcher would pass a mutant that finalizes with the wrong text and still + // leaves the spinner stuck. + const expectedPartialToolMessage = JSON.stringify({ + tool: "newFileCreated", + path: "test/path.txt", + content: testContent, + isOutsideWorkspace: false, + isProtected: false, + }) + mockCline.diffViewProvider.saveChanges.mockRejectedValue(new Error("save failed")) + + await executeWriteFileTool({}) + + expect(mockHandleError).toHaveBeenCalledWith("writing file", expect.any(Error)) + expect(mockCline.finalizePartialToolAsk).toHaveBeenCalledWith(expectedPartialToolMessage) + }) + + it("discards the unapproved preview when the approved write itself fails", async () => { + // saveChanges() releases placeholder ownership only once the write lands, so a rejected + // save leaves this edit owning an empty or half-written new file. The catch has to run + // the discard - the only teardown that removes what this edit created - before reset() + // drops the state the discard reads. + mockCline.diffViewProvider.saveChanges.mockRejectedValue(new Error("save failed")) + mockCline.diffViewProvider.isEditing = true + + await executeWriteFileTool({}) + + const discardOrder = mockCline.diffViewProvider.discardUnapprovedStream.mock.invocationCallOrder[0] + const resetOrder = mockCline.diffViewProvider.reset.mock.invocationCallOrder[0] + expect(mockCline.diffViewProvider.discardUnapprovedStream).toHaveBeenCalledTimes(1) + expect(discardOrder).toBeLessThan(resetOrder) + expect(mockHandleError).toHaveBeenCalledWith("writing file", expect.any(Error)) + }) + + it("still releases the per-task state when the diff view reset fails during the write teardown", async () => { + // reset() sits between the failed write and the release. A reset that rejects must + // neither skip that release nor take over the failure the model is told about. + mockCline.diffViewProvider.saveChanges.mockRejectedValue(new Error("save failed")) + mockCline.diffViewProvider.isEditing = true + mockCline.diffViewProvider.reset.mockRejectedValue(new Error("reset failed")) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + writeToFileTool["getTaskPartialStreamState"](mockCline) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + + try { + await executeWriteFileTool({}) + + expect(mockHandleError).toHaveBeenCalledWith( + "writing file", + expect.objectContaining({ message: "save failed" }), + ) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + } finally { + errorSpy.mockRestore() + writeToFileTool["taskPartialStreamState"].clear() + } + }) + + it("reports a discard that fails during the write teardown", async () => { + // The discard is the last thing standing between an abandoned create and debris on + // disk. If it throws, the user still has to learn the file may be left behind - a + // console line is not a report. + mockCline.diffViewProvider.saveChanges.mockRejectedValue(new Error("save failed")) + mockCline.diffViewProvider.isEditing = true + mockCline.diffViewProvider.discardUnapprovedStream.mockRejectedValue( + new Error("EPERM: operation not permitted"), + ) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + try { + await executeWriteFileTool({}) + + expect(mockCline.say).toHaveBeenCalledWith( + "error", + expect.stringContaining("could not discard the unapproved preview after the failed write"), + ) + expect(mockCline.say).toHaveBeenCalledWith( + "error", + expect.stringContaining("EPERM: operation not permitted"), + ) + // The write failure stays the reported failure; the discard failure is additional. + expect(mockHandleError).toHaveBeenCalledWith( + "writing file", + expect.objectContaining({ message: "save failed" }), + ) + } finally { + errorSpy.mockRestore() + } + }) + + it("finishes the teardown when the discard report itself cannot be delivered", async () => { + // Task.say() throws once the task is aborted. Awaiting the report unguarded let that + // rejection escape the catch, skipping the reset() and the bookkeeping release below - + // so the abort that made the report fail also leaked the state the report described. + mockCline.diffViewProvider.saveChanges.mockRejectedValue(new Error("save failed")) + mockCline.diffViewProvider.isEditing = true + mockCline.diffViewProvider.discardUnapprovedStream.mockRejectedValue( + new Error("EPERM: operation not permitted"), + ) + mockCline.say.mockRejectedValue(new Error("task aborted")) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + writeToFileTool["getTaskPartialStreamState"](mockCline) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + + try { + await executeWriteFileTool({}) + + expect(mockCline.diffViewProvider.reset).toHaveBeenCalledTimes(1) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + expect(mockHandleError).toHaveBeenCalledWith( + "writing file", + expect.objectContaining({ message: "save failed" }), + ) + } finally { + errorSpy.mockRestore() + mockCline.say.mockResolvedValue(undefined) + writeToFileTool["taskPartialStreamState"].clear() + } + }) + }) + + describe("early-exit stream state cleanup", () => { + it("releases this task's stream state when the completed block fails to parse", async () => { + // The streaming deltas registered this task's entry; the finalized block then arrives + // without nativeArgs, so execute() never runs and none of its teardown runs either. + // Without a parse-failure boundary the entry and its TaskAborted listener survive for + // the rest of the task's life, and the diff view keeps unapproved partial content. + await executeWriteFileTool({}, { isPartial: true }) + await executeWriteFileTool({}, { isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + + const block = { + type: "tool_use", + name: "write_to_file", + params: {}, + // No nativeArgs at all: that is what drives BaseTool's parse-failure path, where + // execute() and all of its teardown are skipped. + } as ToolUse<"write_to_file"> + await writeToFileTool.handle(mockCline, block, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + }) + + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + const abortListener = mockCline.once.mock.calls.find( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + )?.[1] + expect(abortListener).toBeInstanceOf(Function) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortListener) + // The stream may have left a diff view open with content that was never approved, so + // the teardown discards it: revertChanges() would SAVE that content to disk. + expect(mockCline.diffViewProvider.discardUnapprovedStream).toHaveBeenCalled() + expect(mockCline.diffViewProvider.revertChanges).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.reset).toHaveBeenCalled() + // Nothing was captured from the stream, so the parse error is still what the user sees. + expect(mockHandleError).toHaveBeenCalledWith("parsing write_to_file args", expect.any(Error)) + }) + + it("releases the per-task stream state when content is missing", async () => { + // The missing-content return sits before execute()'s guarded scope, so it needs its own + // release. The state is seeded first so the assertion proves a release happened rather + // than an empty map. + writeToFileTool["getTaskPartialStreamState"](mockCline as never) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + + const toolUse = { + type: "tool_use", + name: "write_to_file", + params: { path: testFilePath }, + nativeArgs: { path: testFilePath, content: undefined }, + // The fixture's point is a nativeArgs object whose content never arrived, which the + // typed params cannot express - hence the double assertion. + partial: false, + } as unknown as ToolUse<"write_to_file"> + const pushToolResult = vi.fn() + await writeToFileTool.handle(mockCline, toolUse, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult, + }) + + expect(mockCline.sayAndCreateMissingParamError).toHaveBeenCalledWith("write_to_file", "content") + expect(pushToolResult).toHaveBeenCalledWith("Missing param error") + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + const abortListener = mockCline.once.mock.calls.find( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + )?.[1] + expect(abortListener).toBeInstanceOf(Function) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortListener) + // A stream may have opened a diff view for this call; the early return still closes it. + expect(mockCline.diffViewProvider.reset).toHaveBeenCalled() + }) + + it("releases the per-task stream state when the write completes", async () => { + // Seed first: without the seed the map is empty either way and the assertion is vacuous. + writeToFileTool["getTaskPartialStreamState"](mockCline as never) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + + await executeWriteFileTool({}) + + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + const abortListener = mockCline.once.mock.calls.find( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + )?.[1] + expect(abortListener).toBeInstanceOf(Function) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortListener) + expect(mockCline.diffViewProvider.reset).toHaveBeenCalled() + }) + + it("releases the per-task stream state when the write itself fails", async () => { + // The catch path tears down too: a failed write must not leave the entry (and its + // streamFailed guard) attached to the task. + writeToFileTool["getTaskPartialStreamState"](mockCline as never) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + mockCline.diffViewProvider.saveChanges.mockRejectedValue(new Error("save failed")) + + await executeWriteFileTool({}) + + expect(mockHandleError).toHaveBeenCalledWith("writing file", expect.any(Error)) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + const abortListener = mockCline.once.mock.calls.find( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + )?.[1] + expect(abortListener).toBeInstanceOf(Function) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortListener) + }) + it("releases the per-task stream state when a rooignore denial returns early", async () => { + // A partial delta creates the per-task state and registers the abort listener; the + // denial then returns before the cleanup, which used to leave both behind for the + // task's lifetime (a retained streamFailed also suppresses later diff previews). + await executeWriteFileTool({}, { isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + + // A partial delta is not gated by the access check, so the denied call can still be + // holding a preview: open() never consults rooignore. + mockCline.diffViewProvider.isEditing = true + mockCline.diffViewProvider.discardUnapprovedStream.mockClear() + mockCline.diffViewProvider.reset.mockClear() + + await executeWriteFileTool({}, { accessAllowed: false }) + + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + // The exact listener this task registered, not just any function: a mismatched + // off() argument would leave the real listener attached. + const abortListener = mockCline.once.mock.calls.find( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + )?.[1] + expect(abortListener).toBeInstanceOf(Function) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortListener) + // The denied preview must not survive for the next write to reuse. + expect(mockCline.diffViewProvider.discardUnapprovedStream).toHaveBeenCalledTimes(1) + expect(mockCline.diffViewProvider.reset).toHaveBeenCalledTimes(1) + expect(mockCline.diffViewProvider.discardUnapprovedStream.mock.invocationCallOrder[0]).toBeLessThan( + mockCline.diffViewProvider.reset.mock.invocationCallOrder[0], + ) + }) + + it("reports a discard that fails on the rooignore-denial exit", async () => { + // The denial path discards the preview a stream left open before resetting it. When that + // discard throws, the user still has to learn that the preview may be holding content nobody + // approved - a console line is not a report - and the reset below must still run, because it is + // what releases the provider for the next write. + mockCline.diffViewProvider.isEditing = true + mockCline.diffViewProvider.discardUnapprovedStream.mockRejectedValue( + new Error("EPERM: operation not permitted"), + ) + mockCline.diffViewProvider.reset.mockClear() + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + try { + await executeWriteFileTool({}, { accessAllowed: false }) + + // Counted on the exact channel instead of matched with toHaveBeenCalledWith: this exit + // already says "rooignore_error", so only a count of error reports naming the discarded + // preview proves the report happened exactly once. + const discardReports = mockCline.say.mock.calls.filter( + ([type, text]: unknown[]) => + type === "error" && + typeof text === "string" && + text.includes("could not discard the preview for the denied write"), + ) + expect(discardReports).toHaveLength(1) + // The report carries the underlying failure so the user knows what to fix. + expect(discardReports[0][1]).toContain("EPERM: operation not permitted") + expect(mockCline.diffViewProvider.discardUnapprovedStream).toHaveBeenCalledTimes(1) + // A failed report must not skip the teardown below: reset is what releases the diff view. + expect(mockCline.diffViewProvider.reset).toHaveBeenCalledTimes(1) + } finally { + errorSpy.mockRestore() + } + }) + + it("removes the directories a delta adopted before resetting a denied write", async () => { + // The first delta only records the path; a second delta on the same path is what stabilizes + // it, and only then does handlePartial() create the parent directories and hand them to the + // diff view. Empty content keeps that delta from opening a diff view, so the rooignore + // denial reaches the reset with isEditing false, and reset() drops the adopted list without + // touching disk - the directories of a write nobody approved would stay on disk. + const adoptedDirs = ["/mock-workspace/test/nested"] + mockedCreateDirectoriesForFile.mockResolvedValue(adoptedDirs) + await executeWriteFileTool({ content: "" }, { isPartial: true }) + // The first delta stops at the stabilization gate, so nothing is adopted yet. + expect(mockCline.diffViewProvider.adoptCreatedDirectories).not.toHaveBeenCalled() + await executeWriteFileTool({ content: "" }, { isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + // The adopted list the cleanup below removes really holds directories: with only the first + // delta it is empty, the provider removes nothing, and the assertion below cannot fail. + expect(mockCline.diffViewProvider.adoptCreatedDirectories).toHaveBeenCalledWith(adoptedDirs) + expect(mockCline.diffViewProvider.isEditing).toBe(false) + mockCline.diffViewProvider.removeAdoptedDirectories.mockClear() + mockCline.diffViewProvider.reset.mockClear() + mockCline.diffViewProvider.discardUnapprovedStream.mockClear() + + await executeWriteFileTool({}, { accessAllowed: false }) + + expect(mockCline.diffViewProvider.removeAdoptedDirectories).toHaveBeenCalledTimes(1) + // The discard is the session-editing branch's job: with no diff view open the tool must + // not call it, or it reaches past the delta that adopted the directories. + expect(mockCline.diffViewProvider.discardUnapprovedStream).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.reset).toHaveBeenCalledTimes(1) + // The order is the fix: after the reset the adopted list is gone, so removing after it + // would find nothing and leave the directories behind. + expect(mockCline.diffViewProvider.removeAdoptedDirectories.mock.invocationCallOrder[0]).toBeLessThan( + mockCline.diffViewProvider.reset.mock.invocationCallOrder[0], + ) + }) + + it("removes the directories a delta adopted before the validation-rejection release", async () => { + // The same leak on the other exit that skips execute()'s teardown: the stabilized delta + // adopts directories without opening a diff view, then validateToolUse() rejects the + // completed block, so none of the normal cleanup runs. + const adoptedDirs = ["/mock-workspace/test/nested"] + mockedCreateDirectoriesForFile.mockResolvedValue(adoptedDirs) + await executeWriteFileTool({ content: "" }, { isPartial: true }) + await executeWriteFileTool({ content: "" }, { isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + // Same reason as the test above: the adopted list has to hold something before the + // removal can mean anything. + expect(mockCline.diffViewProvider.adoptCreatedDirectories).toHaveBeenCalledWith(adoptedDirs) + expect(mockCline.diffViewProvider.isEditing).toBe(false) + mockCline.diffViewProvider.removeAdoptedDirectories.mockClear() + mockCline.diffViewProvider.reset.mockClear() + + await writeToFileTool.releaseStreamAfterValidationRejection(mockCline) + + expect(mockCline.diffViewProvider.removeAdoptedDirectories).toHaveBeenCalledTimes(1) + expect(mockCline.diffViewProvider.reset).toHaveBeenCalledTimes(1) + expect(mockCline.diffViewProvider.removeAdoptedDirectories.mock.invocationCallOrder[0]).toBeLessThan( + mockCline.diffViewProvider.reset.mock.invocationCallOrder[0], + ) + }) + + it("releases the per-task stream state when a missing parameter returns early", async () => { + await executeWriteFileTool({}, { isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + + await writeToFileTool.execute({ path: "", content: "mock content" }, mockCline, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + }) + + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + expect(mockCline.diffViewProvider.reset).toHaveBeenCalled() + }) + + it("stops before the filesystem probe when the state is released while provider state is in flight", async () => { + // handlePartial() awaits provider.getState() before any side effect. A cancellation + // during that await runs the TaskAborted teardown; the delta already in flight must + // stop there instead of probing, asking and opening a diff view for a dead task. + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + mockCline.diffViewProvider.open.mockClear() + mockCline.providerRef.deref.mockReturnValue({ + getState: vi.fn(async () => { + writeToFileTool.clearTaskState(mockCline) + return {} + }), + }) + + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + + expect(mockCline.ask).toHaveBeenCalledTimes(1) + expect(mockCline.diffViewProvider.open).not.toHaveBeenCalled() + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + }) + + it("stops before the partial ask when the state is released during the filesystem probe", async () => { + // Same teardown, one await later. The first delta pins editType, so clear it to + // take the fileExistsAtPath branch again and abort inside it. + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + mockCline.diffViewProvider.open.mockClear() + mockCline.diffViewProvider.editType = undefined + // mockImplementationOnce: executeWriteFileTool re-arms the default resolved value + // on every call, so a plain mockImplementation would be overwritten. + mockedFileExistsAtPath.mockImplementationOnce(async () => { + writeToFileTool.clearTaskState(mockCline) + return false + }) + + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + + expect(mockCline.ask).toHaveBeenCalledTimes(1) + expect(mockCline.diffViewProvider.open).not.toHaveBeenCalled() + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + }) + + it("removes the directories it created when the stream is released during their creation", async () => { + // The delta creates the parent directories right after the probe and hands them to the + // diff view's cleanup state. A cancellation landing inside that creation leaves + // directories no teardown will visit: open() never ran, so the discard and the revert + // have no session to clean, and reset() drops the recorded list without touching disk. + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + mockCline.diffViewProvider.editType = undefined + mockCline.diffViewProvider.adoptCreatedDirectories.mockClear() + mockCline.diffViewProvider.removeAdoptedDirectories.mockClear() + mockCline.diffViewProvider.open.mockClear() + const created = ["/mock-workspace/new-file/nested"] + // mockImplementationOnce: executeWriteFileTool re-arms the default resolved value + // on every call, so a plain mockImplementation would be overwritten. + mockedCreateDirectoriesForFile.mockImplementationOnce(async () => { + writeToFileTool.clearTaskState(mockCline) + return created + }) + + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + + expect(mockCline.diffViewProvider.adoptCreatedDirectories).toHaveBeenCalledWith(created) + expect(mockCline.diffViewProvider.removeAdoptedDirectories).toHaveBeenCalledTimes(1) + expect(mockCline.diffViewProvider.open).not.toHaveBeenCalled() + }) + + it("stops before touching the diff view when the stream state is released during an in-flight ask", async () => { + // A cancellation while task.ask() is in flight runs the TaskAborted teardown. The + // delta that was already in flight must not then re-open the diff view for a task + // the user cancelled - that resurrects the state the teardown just released. + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + mockCline.diffViewProvider.open.mockClear() + mockCline.diffViewProvider.update.mockClear() + mockCline.ask.mockImplementation(async () => { + writeToFileTool.clearTaskState(mockCline) + }) + + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + + expect(mockCline.diffViewProvider.open).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.update).not.toHaveBeenCalled() + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + }) + + it("stops before updating the diff view when the task is cancelled while open() is in flight", async () => { + // open() is the first provider await after the partial ask. If TaskAborted lands + // while it is in flight, the teardown has already released this task's stream + // state (and may have reverted or closed this very view), so the delta that is + // already in flight must not stream partial content into a cancelled task's view. + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + expect(mockCline.diffViewProvider.open).toHaveBeenCalledTimes(1) + mockCline.diffViewProvider.open.mockClear() + mockCline.diffViewProvider.update.mockClear() + mockCline.diffViewProvider.discardUnapprovedStream.mockClear() + mockCline.diffViewProvider.reset.mockClear() + mockCline.diffViewProvider.open.mockImplementationOnce(async () => { + // open() had already marked the session as editing when the abort landed - which + // is exactly why the abort cleanup, checking isEditing earlier, could not close it. + mockCline.diffViewProvider.isEditing = true + writeToFileTool.clearTaskState(mockCline) + }) + + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + + expect(mockCline.diffViewProvider.open).toHaveBeenCalledTimes(1) + expect(mockCline.diffViewProvider.update).not.toHaveBeenCalled() + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + // The view and the placeholder open() wrote are this delta's to close: the abort + // cleanup had already run, so nobody else would. + expect(mockCline.diffViewProvider.discardUnapprovedStream).toHaveBeenCalledTimes(1) + expect(mockCline.diffViewProvider.reset).toHaveBeenCalledTimes(1) + }) + it("releases the per-task stream state when the user rejects the diff-view approval", async () => { + // A partial delta registers the entry and the TaskAborted listener. The denial then + // returns from inside the try block, skipping the success-path teardown, so both stay + // attached for the rest of the task's life (and a retained streamFailed would keep + // suppressing this task's later diff previews). + await executeWriteFileTool({}, { isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + mockAskApproval.mockResolvedValue(false) + await executeWriteFileTool({}) + expect(mockCline.diffViewProvider.revertChanges).toHaveBeenCalledTimes(1) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + // The exact listener this task registered, not just any function: a mismatched + // off() argument would leave the real listener attached. + const abortListener = mockCline.once.mock.calls.find( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + )?.[1] + expect(abortListener).toBeInstanceOf(Function) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortListener) + }) + it("releases the per-task stream state when the prevent-focus-disruption approval is rejected", async () => { + // The experiment branch asks for approval without ever opening a diff view, so the + // only teardown for this call is the one at the end of the try block - which the + // denial return skips. + await executeWriteFileTool({}, { isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + mockCline.providerRef.deref.mockReturnValue({ + getState: vi.fn().mockResolvedValue({ + diagnosticsEnabled: true, + writeDelayMs: 1000, + experiments: { preventFocusDisruption: true }, + }), + }) + mockCline.diffViewProvider.saveDirectly = vi.fn().mockResolvedValue(undefined) + mockAskApproval.mockResolvedValue(false) + await executeWriteFileTool({}) + expect(mockCline.diffViewProvider.saveDirectly).not.toHaveBeenCalled() + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + // The exact listener this task registered, not just any function: a mismatched + // off() argument would leave the real listener attached. + const abortListener = mockCline.once.mock.calls.find( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + )?.[1] + expect(abortListener).toBeInstanceOf(Function) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortListener) + }) + it("releases the per-task stream state when prevent-focus-disruption skips the partial preview", async () => { + // The first delta only pins the path, so the entry is still live after it (the stream + // is in flight). The second delta reaches the experiment check: handlePartial() then + // returns without ever showing a preview, and nothing else would ever release the + // entry or detach the TaskAborted listener for this task. + mockCline.providerRef.deref.mockReturnValue({ + getState: vi.fn().mockResolvedValue({ experiments: { preventFocusDisruption: true } }), + }) + + await executeWriteFileTool({}, { isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + + await executeWriteFileTool({}, { isPartial: true }) + + expect(mockCline.ask).not.toHaveBeenCalled() + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + // The exact listener this task registered, not just any function: a mismatched + // off() argument would leave the real listener attached. + const abortListener = mockCline.once.mock.calls.find( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + )?.[1] + expect(abortListener).toBeInstanceOf(Function) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortListener) + }) + + it("stops the partial delta when the task aborts while directory creation is in flight", async () => { + // handlePartial() awaits createDirectoriesForFile() for a new file, then asks and streams + // the diff view. An abandonment that lands during that await has already released this + // task's stream state; without a re-check after the await the delta keeps going and puts a + // partial tool ask and a diff-view update on screen for a task that no longer exists. + let abortCleanup: (() => void) | undefined + mockCline.once.mockImplementation((event: RooCodeEventName, listener: () => void) => { + if (event === RooCodeEventName.TaskAborted) { + abortCleanup = listener + } + return mockCline + }) + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + + let releaseGate: (() => void) | undefined + const gate = new Promise((resolve) => { + releaseGate = resolve + }) + mockedCreateDirectoriesForFile.mockImplementationOnce(() => gate.then(() => [])) + const streaming = executeWriteFileTool({}, { fileExists: false, isPartial: true }) + await new Promise((resolve) => setImmediate(resolve)) + expect(mockedCreateDirectoriesForFile).toHaveBeenCalledTimes(1) + + abortCleanup?.() + releaseGate?.() + await streaming + + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + expect(mockCline.ask).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.open).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.update).not.toHaveBeenCalled() + }) + + it("releases the per-task stream state when provider state rejects during a partial delta", async () => { + // handlePartial() registers the entry and the TaskAborted listener, then awaits + // provider.getState(). A rejection there never reaches the diff view or execute(), so + // nothing else releases what the registration acquired. The error still has to surface, + // so the boundary rethrows and BaseTool.handle() reports it once. + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + mockCline.providerRef.deref.mockReturnValue({ + getState: vi.fn().mockRejectedValue(new Error("provider state unavailable")), + }) + + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + const abortListener = mockCline.once.mock.calls.find( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + )?.[1] + expect(abortListener).toBeInstanceOf(Function) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortListener) + expect(mockHandleError).toHaveBeenCalledWith( + "handling partial write_to_file", + expect.objectContaining({ message: "provider state unavailable" }), + ) + }) + + it("reports a discard failure without replacing the error the delta produced", async () => { + // The delta failed in the pre-streaming setup, and the discard of the preview an + // earlier delta left open failed too. Two failures, two channels: BaseTool.handle() + // must still report THIS delta's error - wrapping it would change the failure the + // caller sees - while the discard failure (debris still on disk) gets its own report. + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + mockCline.providerRef.deref.mockReturnValue({ + getState: vi.fn().mockRejectedValue(new Error("provider state unavailable")), + }) + mockCline.diffViewProvider.isEditing = true + mockCline.diffViewProvider.discardUnapprovedStream.mockRejectedValue( + new Error("EPERM: operation not permitted"), + ) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + try { + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + + expect(mockCline.say).toHaveBeenCalledWith( + "error", + expect.stringContaining("could not discard the unapproved preview after the failed stream"), + ) + expect(mockCline.say).toHaveBeenCalledWith( + "error", + expect.stringContaining("EPERM: operation not permitted"), + ) + const reported = mockHandleError.mock.calls.find( + ([context]) => context === "handling partial write_to_file", + )?.[1] as Error + expect(reported.message).toBe("provider state unavailable") + expect(reported.name).toBe("Error") + } finally { + errorSpy.mockRestore() + } + }) + + it("still reports the delta's own error when the discard report cannot be delivered", async () => { + // Same guard on the streaming side. The exception this delta produced is what + // BaseTool.handle() reports; an undeliverable report must not take its place, or the + // caller sees an abort where a provider failure happened. + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + mockCline.providerRef.deref.mockReturnValue({ + getState: vi.fn().mockRejectedValue(new Error("provider state unavailable")), + }) + mockCline.diffViewProvider.isEditing = true + mockCline.diffViewProvider.discardUnapprovedStream.mockRejectedValue( + new Error("EPERM: operation not permitted"), + ) + mockCline.say.mockRejectedValue(new Error("task aborted")) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + try { + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + + const reported = mockHandleError.mock.calls.find( + ([context]) => context === "handling partial write_to_file", + )?.[1] as Error + expect(reported.message).toBe("provider state unavailable") + expect(mockCline.say).toHaveBeenCalledWith( + "error", + expect.stringContaining("could not discard the unapproved preview after the failed stream"), + ) + } finally { + errorSpy.mockRestore() + mockCline.say.mockResolvedValue(undefined) + } + }) + + it("leaves a replacement stream state alone when the delta it replaced resumes", async () => { + // Liveness is object identity, not key presence: a test that only clears the entry + // passes for an implementation that checks the key. Here a new stream takes over the + // same task key while the old delta awaits the filesystem, so the resumed delta must + // neither continue its side effects nor write through the replacement's entry. + let releaseGate: (() => void) | undefined + const gate = new Promise((resolve) => { + releaseGate = resolve + }) + // The path must be stabilized by an earlier delta before handlePartial() touches the + // filesystem, which is also what registers this task's entry. + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + mockedCreateDirectoriesForFile.mockImplementationOnce(() => gate.then(() => [])) + const streaming = executeWriteFileTool({}, { fileExists: false, isPartial: true }) + await new Promise((resolve) => setImmediate(resolve)) + expect(mockedCreateDirectoriesForFile).toHaveBeenCalledTimes(1) + + const key = writeToFileTool["getPartialStreamFailureKey"](mockCline as never) as string + writeToFileTool["resetTaskPartialState"](mockCline as never) + const replacement = writeToFileTool["getTaskPartialStreamState"](mockCline as never) + replacement.streamFailed = false + releaseGate?.() + await streaming + + expect(mockCline.ask).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.open).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.update).not.toHaveBeenCalled() + expect(writeToFileTool["taskPartialStreamState"].get(key)).toBe(replacement) + expect(replacement.streamFailed).toBe(false) + expect(replacement.streamError).toBeUndefined() + }) + + it("releases the stream state when the setup before the write boundary rejects", async () => { + // execute() creates the parent directories for a new file before its write boundary + // begins. An EACCES there used to escape execute() entirely - BaseTool.handle() only + // reports it - so the map kept this task's entry and its abort listener alive, and a + // later write could inherit a stale stream state. + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + const abortListener = mockCline.once.mock.calls.find( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + )?.[1] + expect(abortListener).toBeInstanceOf(Function) + mockedCreateDirectoriesForFile.mockRejectedValueOnce(new Error("EACCES: permission denied")) + + // The failure itself keeps travelling the path it always took: the boundary releases + // and rethrows, so the caller still sees the original error and nothing reports it twice. + await expect(executeWriteFileTool({})).rejects.toThrow("EACCES: permission denied") + + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortListener) + expect(mockHandleError).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.reset).not.toHaveBeenCalled() + }) + + it("reports the captured stream error once when the finalized block fails to parse", async () => { + // A streaming delta already hit a fatal filesystem error; the finalized block then fails + // to parse. The stream error is what the user can act on, so it takes the report slot and + // the incidental parse error is suppressed - reporting both would show two bubbles for one + // failure, reporting only the parse error would drop the actionable one. The error is + // induced through open() rather than assigned, so what gets reported is what production + // recorded on the entry. + const streamError = new Error("EROFS: read-only file system, open '/ro/test.py'") + mockCline.diffViewProvider.open.mockImplementation(async () => { + mockCline.diffViewProvider.isEditing = true + throw streamError + }) + await executeWriteFileTool({}, { isPartial: true }) + await executeWriteFileTool({}, { isPartial: true }) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + mockHandleError.mockClear() + + const block = { + type: "tool_use", + name: "write_to_file", + params: {}, + // No nativeArgs at all: that is what drives BaseTool's parse-failure path. + partial: false, + } as ToolUse<"write_to_file"> + await writeToFileTool.handle(mockCline, block, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + }) + + expect(mockHandleError).toHaveBeenCalledTimes(1) + expect(mockHandleError).toHaveBeenCalledWith("writing file", streamError) + expect(mockHandleError).not.toHaveBeenCalledWith("parsing write_to_file args", expect.any(Error)) + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + }) + + it("reports a failed rollback as the cleanup failure when the finalized block fails to parse", async () => { + // The rollback is what keeps unapproved streamed content off disk. When it fails, the + // debris is the more actionable failure: it takes the report slot with the stream error + // kept behind it as the cause, instead of being logged and continued past. + const streamError = new Error("EACCES: permission denied, open '/ro/test.py'") + mockCline.diffViewProvider.open.mockImplementation(async () => { + mockCline.diffViewProvider.isEditing = true + throw streamError + }) + await executeWriteFileTool({}, { isPartial: true }) + await executeWriteFileTool({}, { isPartial: true }) + mockCline.diffViewProvider.discardUnapprovedStream.mockRejectedValue( + new Error("EACCES: could not remove the directory created for this write"), + ) + mockHandleError.mockClear() + + const block = { + type: "tool_use", + name: "write_to_file", + params: {}, + partial: false, + } as ToolUse<"write_to_file"> + await writeToFileTool.handle(mockCline, block, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + }) + + expect(mockHandleError).toHaveBeenCalledTimes(1) + expect(mockHandleError).toHaveBeenCalledWith( + "writing file", + expect.objectContaining({ message: expect.stringContaining("rollback failed") }), + ) + expect(mockHandleError.mock.calls[0]?.[1]).toHaveProperty("cause", streamError) + expect(mockHandleError).not.toHaveBeenCalledWith("parsing write_to_file args", expect.any(Error)) + }) + + it("discards the unapproved preview and records the stream error when open() rejects", async () => { + // open() sits inside the cleanup-owned boundary: the preview is discarded (never saved), + // the view reset, the error recorded on this call's entry, and the error rethrown so + // BaseTool.handle() reports it exactly once. + const failure = new Error("EPERM: could not open the diff editor") + mockCline.diffViewProvider.open.mockImplementation(async () => { + mockCline.diffViewProvider.isEditing = true + throw failure + }) + + await executeWriteFileTool({}, { isPartial: true }) + await executeWriteFileTool({}, { isPartial: true }) + + expect(mockHandleError).toHaveBeenCalledTimes(1) + expect(mockHandleError).toHaveBeenCalledWith("handling partial write_to_file", failure) + expect(mockCline.diffViewProvider.discardUnapprovedStream).toHaveBeenCalledTimes(1) + // revertChanges() SAVES: an unapproved preview must never be routed through it. + expect(mockCline.diffViewProvider.revertChanges).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.reset).toHaveBeenCalledTimes(1) + const retained = [...writeToFileTool["taskPartialStreamState"].values()][0] + expect(writeToFileTool["taskPartialStreamState"].size).toBe(1) + expect(retained.streamFailed).toBe(true) + expect(retained.streamError).toBe(failure) + // The listener goes with whichever teardown ends the call, not with this delta: the + // retained entry is what suppresses the rest of the stream. + expect(mockCline.off).not.toHaveBeenCalled() + }) + + it("suppresses the rest of the stream once the diff view has failed for this call", async () => { + const failure = new Error("EPERM: could not open the diff editor") + mockCline.diffViewProvider.open.mockImplementation(async () => { + mockCline.diffViewProvider.isEditing = true + throw failure + }) + await executeWriteFileTool({}, { isPartial: true }) + await executeWriteFileTool({}, { isPartial: true }) + + mockCline.diffViewProvider.open.mockClear() + mockCline.diffViewProvider.update.mockClear() + mockCline.diffViewProvider.discardUnapprovedStream.mockClear() + mockCline.ask.mockClear() + mockHandleError.mockClear() + + // A third delta for the same call. Retrying would re-open the diff editor that just + // failed and re-ask for a call that already reported its error. + await executeWriteFileTool({}, { isPartial: true }) + + expect(mockCline.diffViewProvider.open).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.update).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.discardUnapprovedStream).not.toHaveBeenCalled() + expect(mockCline.ask).not.toHaveBeenCalled() + expect(mockHandleError).not.toHaveBeenCalled() + }) + + it("discards the unapproved preview and records the stream error when update() rejects", async () => { + // The other half of the boundary: open() succeeded, so the view holds partial content, + // and update() is what fails. + const failure = new Error("EPERM: could not stream into the diff editor") + mockCline.diffViewProvider.open.mockImplementation(async () => { + mockCline.diffViewProvider.isEditing = true + }) + mockCline.diffViewProvider.update.mockRejectedValue(failure) + + await executeWriteFileTool({}, { isPartial: true }) + await executeWriteFileTool({}, { isPartial: true }) + + expect(mockHandleError).toHaveBeenCalledTimes(1) + expect(mockHandleError).toHaveBeenCalledWith("handling partial write_to_file", failure) + expect(mockCline.diffViewProvider.discardUnapprovedStream).toHaveBeenCalledTimes(1) + expect(mockCline.diffViewProvider.revertChanges).not.toHaveBeenCalled() + expect(mockCline.diffViewProvider.reset).toHaveBeenCalledTimes(1) + const retained = [...writeToFileTool["taskPartialStreamState"].values()][0] + expect(retained.streamFailed).toBe(true) + expect(retained.streamError).toBe(failure) + }) + + it("releases the stream state and deregisters the abort listener when the failed delta's block never completes", async () => { + const failure = new Error("EPERM: could not open the diff editor") + mockCline.diffViewProvider.open.mockImplementation(async () => { + mockCline.diffViewProvider.isEditing = true + throw failure + }) + await executeWriteFileTool({}, { isPartial: true }) + await executeWriteFileTool({}, { isPartial: true }) + mockHandleError.mockClear() + + // The stream died mid-parameters, so the finalized block never parses and execute() never + // runs: the parse-failure boundary is what ends the call and releases the entry. + const block = { + type: "tool_use", + name: "write_to_file", + params: {}, + partial: false, + } as ToolUse<"write_to_file"> + await writeToFileTool.handle(mockCline, block, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + }) + + expect(writeToFileTool["taskPartialStreamState"].size).toBe(0) + const abortListener = mockCline.once.mock.calls.find( + ([event]: unknown[]) => event === RooCodeEventName.TaskAborted, + )?.[1] + expect(abortListener).toBeInstanceOf(Function) + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortListener) + // The induced stream error is the actionable one, reported once; the incidental parse + // error is suppressed. + expect(mockHandleError).toHaveBeenCalledTimes(1) + expect(mockHandleError).toHaveBeenCalledWith("writing file", failure) + expect(mockHandleError).not.toHaveBeenCalledWith("parsing write_to_file args", expect.any(Error)) + }) }) describe("user interaction", () => { diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 07772816b4..70f2aa6503 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -60,6 +60,7 @@ import { } from "@roo-code/types" import { RateLimitClock, createRateLimitClock } from "../task/RateLimitClock" import { TaskRegistry } from "../task/TaskRegistry" +import { writeToFileTool } from "../tools/WriteToFileTool" import { TaskScheduler } from "../task/TaskScheduler" import { getEffectiveTaskApiConfiguration, @@ -638,6 +639,11 @@ export class ClineProvider this.taskEventListeners.delete(task) } + // Dispose removes all listeners, so the tool's per-task state must be released + // while the task is still here; otherwise the singleton keeps the disposed task + // and its diff-view provider. + writeToFileTool.clearTaskState(task) + try { await task.dispose() } catch (error) { diff --git a/src/integrations/editor/DiffViewProvider.ts b/src/integrations/editor/DiffViewProvider.ts index bb3368f063..a6c788723a 100644 --- a/src/integrations/editor/DiffViewProvider.ts +++ b/src/integrations/editor/DiffViewProvider.ts @@ -33,6 +33,14 @@ export class DiffViewProvider { isEditing = false originalContent: string | undefined private createdDirs: string[] = [] + /** + * Absolute path of the empty placeholder open() created for a new-file edit, or undefined + * when no placeholder is outstanding. reset() does NOT clear relPath, and saveDirectly() + * sets relPath for a file that was written with approval, so relPath alone cannot tell a + * caller that a file on disk belongs to the abandoned edit. Only this field justifies an + * unlink. + */ + private placeholderPath: string | undefined private documentWasOpen = false // Tracks whether the target file's tab was pinned before the diff session. // Closing the tab to open the diff drops VS Code's pin state, so we restore @@ -126,12 +134,18 @@ export class DiffViewProvider { } // For new files, create any necessary directories and keep track of new - // directories to delete if the user denies the operation. - this.createdDirs = await createDirectoriesForFile(absolutePath) + // directories to delete if the user denies the operation. Merged rather than + // assigned: handlePartial() may have created them first and recorded them, and + // overwriting the list would hand those directories back to the leak. + this.adoptCreatedDirectories(await createDirectoriesForFile(absolutePath)) // Make sure the file exists before we open it. if (!fileExists) { await fs.writeFile(absolutePath, "") + // From here until the placeholder is removed, THIS edit owns that path. Set after + // the write succeeds, so a failed write does not claim ownership of a file we did + // not create. + this.placeholderPath = absolutePath } // If the file was already open, close it (must happen after showing the @@ -343,6 +357,12 @@ export class DiffViewProvider { await updatedDocument.save() } + // Ownership is released only once the approved write has landed. A save that rejects leaves + // the placeholder as open() created it, and the caller's teardown still has to remove it: + // clearing it here would tell discardUnapprovedStream() that nothing this edit created is + // left to clean, and an empty or half-written new file would survive the failed write. + this.placeholderPath = undefined + // Stop tracking touches and cancel any pending scroll-to-diff before any // programmatic editor activation below. this.disposeActiveEditorListener() @@ -513,6 +533,172 @@ export class DiffViewProvider { return JSON.stringify(result) } + /** + * Release a diff view whose content was never approved: an abandoned partial stream, + * a failed stream, or a write the user was never asked to approve (rooignore denial, + * validation failure). Deliberately NOT the same as revertChanges(), because + * revertChanges() SAVES and neither abandoned case may write to disk: + * + * - create: its new-file branch saves a dirty buffer as-is before deleting the file, + * which writes partial model output the task never approved - and leaves it on disk + * if the delete then fails. + * - modify: its branch restores the original content and saves it. That is a write to + * a file the user never approved, and for a rooignore-denial path a write the + * policy forbids outright. + * + * Here the target file is never written. A create buffer is emptied FIRST, so the + * only bytes that can ever reach the placeholder are none and the tab is clean enough + * to close without a prompt; a modify buffer is restored to the content already on + * disk in memory only, so the file itself is untouched and the tab (now showing the + * original content, still marked dirty) is left for the user rather than force-closed + * over any edits they may have typed into the preview. The placeholder plus the + * directories this edit created are removed either way. + */ + async discardUnapprovedStream(): Promise { + if (!this.relPath) { + return + } + + const absolutePath = path.resolve(this.cwd, this.relPath) + // Snapshot the directories this edit created BEFORE the editor work below, and + // clear the field so nothing else can act on them twice. If an await below + // rejects, the caller runs reset(), which drops relPath/createdDirs - this is + // then the only chance to remove what the abandoned edit left on disk. + const createdDirs = this.createdDirs + this.createdDirs = [] + // Same reason for the placeholder: an await below may reject and the caller then runs + // reset(), which drops this field too. + const placeholderPath = this.placeholderPath + this.placeholderPath = undefined + + let editorFailure: unknown + // Whether the buffer still holds unapproved content once the editor work is over. A dirty + // buffer is what makes the placeholder unsafe to remove: its tab is still open, and an + // unlink would leave that tab pointing at a path with no file behind it. + let bufferStillDirty = false + // open() creates the directories and the empty placeholder BEFORE it assigns + // activeDiffEditor (openDiffEditor() can reject on its 10s timeout or a failed + // vscode.diff call), so an abandoned create can leave artifacts on disk with no + // editor at all. Only the buffer and tab work needs the editor; the artifact + // cleanup below runs either way. + if (this.activeDiffEditor) { + const document = this.activeDiffEditor.document + try { + this.disposeActiveEditorListener() + this.cancelDeferredScroll() + + if (document.isDirty) { + // Restore the buffer to what this edit started from - the empty placeholder for a + // create, the content already on disk for a modify. A failed applyEdit leaves the + // unapproved streamed content in the buffer, and saving then would persist exactly + // what this method exists to discard, so the save is conditional and the failure is + // surfaced to the caller as a rollback hazard. + const edit = new vscode.WorkspaceEdit() + const fullRange = new vscode.Range( + document.positionAt(0), + document.positionAt(document.getText().length), + ) + const restoredContent = + this.editType === "modify" ? this.stripAllBOMs(this.originalContent ?? "") : "" + edit.replace(document.uri, fullRange, restoredContent) + const applied = await vscode.workspace.applyEdit(edit) + if (!applied) { + editorFailure = new Error( + `Could not restore the diff editor buffer for ${this.relPath}; it may still hold unapproved content.`, + ) + } else if (this.editType !== "modify") { + // Only the emptied placeholder is ever saved: this edit created that file, and + // closing a dirty tab would prompt. A modify is never saved here - a modify's + // restore is an in-memory revert, and saving it would write to a file the user + // never approved (see the method comment). + const saved = await document.save() + if (!saved) { + // A save the editor did not perform leaves the unapproved content in the buffer. + // Reading it as success would let the cleanup below delete the placeholder under a + // dirty tab and still report a restored preview - the next Ctrl+S would then + // recreate the file with exactly what this method exists to discard. + editorFailure = new Error( + `Could not save the restored diff editor buffer for ${this.relPath}; it may still hold unapproved content.`, + ) + } + } + } + + // Close the diff tabs only now: a vscode.diff tab is dirty while its modified side + // holds the streamed content, and closeAllDiffViews() deliberately skips dirty tabs to + // avoid a save prompt. Closing before the restore above left that tab open over a file + // this method is about to unlink. + await this.closeAllDiffViews() + await this.closeFileTab(absolutePath) + } catch (error) { + // Do NOT stop here: the placeholder and the created directories still have + // to go. The original failure is re-thrown once the artifacts are dealt with, + // so the caller still reports the rollback hazard instead of a silent success. + editorFailure = error + } + bufferStillDirty = document.isDirty + } + + // Set when the editor work failed while the buffer still held unapproved content: the + // open tab is then the only thing that can still write that content to disk, so the + // rollback leaves the artifacts alone rather than orphaning the editor. + const keepArtifacts = editorFailure !== undefined && bufferStillDirty + + let cleanupFailure: unknown + try { + // Only a placeholder THIS edit created may be removed. relPath survives reset(), and + // saveDirectly() sets it for a file that was written with approval, so unlinking + // absolutePath unconditionally can delete content the user approved. + // + // A rollback that failed while the buffer was still dirty is a second reason to leave + // it alone: the tab is open, so unlinking now would leave it pointing at a path with no + // file behind it, and the next Ctrl+S would recreate the file with exactly the + // unapproved content this method exists to discard. Keeping the placeholder costs a + // stray empty file; the alternative throws the user's save into a phantom. When the + // buffer ended clean - or there was never an editor at all, the open()-failed-before- + // the-editor case - the artifact really is debris and goes. + if (placeholderPath && !keepArtifacts) { + await fs.unlink(placeholderPath).catch((error: unknown) => { + // open() creates the placeholder; if it is already gone there is nothing + // left to remove and the discard did its job. + if ((error as NodeJS.ErrnoException)?.code !== "ENOENT") { + throw error + } + }) + } + + // Remove only the directories this edit created, in reverse order - and not at all + // when the placeholder was kept, since the kept file still lives inside the deepest + // of them and rmdir would fail on a non-empty directory anyway. + if (!keepArtifacts) { + for (let i = createdDirs.length - 1; i >= 0; i--) { + await fs.rmdir(createdDirs[i]).catch((error: unknown) => { + if ((error as NodeJS.ErrnoException)?.code !== "ENOENT") { + throw error + } + }) + } + } + } catch (error) { + cleanupFailure = error + console.error("Error removing abandoned write_to_file artifacts:", error) + } + + if (editorFailure) { + const primary = editorFailure instanceof Error ? editorFailure : new Error(String(editorFailure)) + if (keepArtifacts) { + // Same error object and type, with the state of the disk appended: the caller + // reports this text, and "the rollback failed" is only actionable once the reader + // knows what was left where. + primary.message = `${primary.message} The placeholder file and the directories this edit created were kept on disk because the editor is still open on them.` + } + throw primary + } + if (cleanupFailure) { + throw cleanupFailure instanceof Error ? cleanupFailure : new Error(String(cleanupFailure)) + } + } + async revertChanges(): Promise { if (!this.relPath || !this.activeDiffEditor) { return @@ -537,6 +723,8 @@ export class DiffViewProvider { // opened tab before deleting it from disk. await this.closeFileTab(absolutePath) await fs.unlink(absolutePath) + // The placeholder this edit owned is gone; no later cleanup may unlink that path again. + this.placeholderPath = undefined // Remove only the directories we created, in reverse order. for (let i = this.createdDirs.length - 1; i >= 0; i--) { @@ -1099,6 +1287,52 @@ export class DiffViewProvider { return result } + /** + * Record directories an edit created before the diff view exists. + * + * handlePartial() creates the parent directories early so a later open() cannot hit ENOENT, + * and open() can only record the directories IT creates - which is none, once that earlier + * call has already made them. Every teardown reads createdDirs (the discard and the revert + * both remove it) while reset() merely drops the list, so an unrecorded creation is + * invisible to all of them and survives on disk. Merging rather than assigning keeps + * whichever call made the directories first still accounted for. + */ + adoptCreatedDirectories(directories: string[]): void { + for (const directory of directories) { + if (!this.createdDirs.includes(directory)) { + this.createdDirs.push(directory) + } + } + } + + /** + * Remove the directories this edit recorded but the diff view never took over. + * + * A delta that created parent directories and then stopped - a cancellation landing during + * the creation, or a setup failure before open() - leaves directories no teardown will + * visit: the discard and the revert only run for a session that reached the editor, and + * reset() drops the list without touching the disk. Deepest first; a directory that is + * already gone or not empty is not this edit's to force away. + */ + async removeAdoptedDirectories(): Promise { + const directories = this.createdDirs + this.createdDirs = [] + for (let i = directories.length - 1; i >= 0; i--) { + try { + await fs.rmdir(directories[i]) + } catch (error: unknown) { + // A directory that is already gone, or that something else now lives in, is not + // this edit's to force away. Anything else - a permission failure, an I/O error - + // is a directory left behind that nothing still tracks, so it is reported rather + // than swallowed. The remaining directories are attempted either way. + const code = (error as NodeJS.ErrnoException)?.code + if (code !== "ENOENT" && code !== "ENOTEMPTY") { + console.error("Error removing a directory this edit created:", directories[i], error) + } + } + } + } + async reset(): Promise { // Dispose touch listeners and cancel any pending deferred scroll BEFORE any // async editor manipulation. closeAllDiffViews() awaits tab-close operations, @@ -1113,6 +1347,7 @@ export class DiffViewProvider { this.isEditing = false this.originalContent = undefined this.createdDirs = [] + this.placeholderPath = undefined this.documentWasOpen = false this.documentWasPinned = false this.activeDiffEditor = undefined diff --git a/src/integrations/editor/__tests__/DiffViewProvider.spec.ts b/src/integrations/editor/__tests__/DiffViewProvider.spec.ts index 00b3dcaf7a..036e9aeca9 100644 --- a/src/integrations/editor/__tests__/DiffViewProvider.spec.ts +++ b/src/integrations/editor/__tests__/DiffViewProvider.spec.ts @@ -1,3 +1,4 @@ +import * as fs from "fs/promises" import { DiffViewProvider, DIFF_VIEW_URI_SCHEME, DIFF_VIEW_LABEL_CHANGES } from "../DiffViewProvider" import * as vscode from "vscode" import * as path from "path" @@ -15,6 +16,10 @@ vi.mock("delay", () => ({ vi.mock("fs/promises", () => ({ readFile: vi.fn().mockResolvedValue("file content"), writeFile: vi.fn().mockResolvedValue(undefined), + // The abandoned-stream discard removes the placeholder file and the directories it + // created, so the discard path needs these two. + unlink: vi.fn().mockResolvedValue(undefined), + rmdir: vi.fn().mockResolvedValue(undefined), access: vi.fn().mockResolvedValue(undefined), })) @@ -36,7 +41,7 @@ vi.mock("vscode", () => ({ onDidOpenTextDocument: vi.fn(() => ({ dispose: vi.fn() })), openTextDocument: vi.fn().mockResolvedValue({ isDirty: false, - save: vi.fn().mockResolvedValue(undefined), + save: vi.fn().mockResolvedValue(true), }), textDocuments: [], fs: { @@ -870,7 +875,7 @@ describe("DiffViewProvider", () => { document: { getText: vi.fn().mockReturnValue("new content"), isDirty: false, - save: vi.fn().mockResolvedValue(undefined), + save: vi.fn().mockResolvedValue(true), }, } ;(diffViewProvider as any).preDiagnostics = [] @@ -1015,7 +1020,7 @@ describe("DiffViewProvider", () => { document: { getText: vi.fn().mockReturnValue("content"), isDirty: false, - save: vi.fn().mockResolvedValue(undefined), + save: vi.fn().mockResolvedValue(true), }, } @@ -1040,7 +1045,7 @@ describe("DiffViewProvider", () => { document: { getText: vi.fn().mockReturnValue("content"), isDirty: false, - save: vi.fn().mockResolvedValue(undefined), + save: vi.fn().mockResolvedValue(true), }, } @@ -1060,7 +1065,7 @@ describe("DiffViewProvider", () => { uri: { fsPath: `${mockCwd}/race.ts`, scheme: "file" }, getText: vi.fn().mockReturnValue("a\nCHANGED\nc\nd\n"), isDirty: false, - save: vi.fn().mockResolvedValue(undefined), + save: vi.fn().mockResolvedValue(true), lineCount: 5, lineAt: vi.fn().mockReturnValue({ text: "" }), }, @@ -1108,7 +1113,7 @@ describe("DiffViewProvider", () => { uri: { fsPath: `${mockCwd}/test.txt` }, getText: vi.fn().mockReturnValue("modified"), isDirty: false, - save: vi.fn().mockResolvedValue(undefined), + save: vi.fn().mockResolvedValue(true), positionAt: vi.fn().mockReturnValue({ line: 0, character: 0 }), }, } @@ -1132,7 +1137,7 @@ describe("DiffViewProvider", () => { uri: { fsPath: mockTargetPath }, getText: vi.fn().mockReturnValue("content"), isDirty: false, - save: vi.fn().mockResolvedValue(undefined), + save: vi.fn().mockResolvedValue(true), positionAt: vi.fn().mockReturnValue({ line: 0, character: 0 }), }, }) @@ -1258,7 +1263,7 @@ describe("DiffViewProvider", () => { uri: { fsPath: mockTargetPath }, getText: vi.fn().mockReturnValue("content"), isDirty: false, - save: vi.fn().mockResolvedValue(undefined), + save: vi.fn().mockResolvedValue(true), positionAt: vi.fn().mockReturnValue({ line: 0, character: 0 }), }, }) @@ -1664,7 +1669,7 @@ describe("DiffViewProvider", () => { uri: { fsPath: mockTargetPath }, getText: vi.fn().mockReturnValue("content"), isDirty: false, - save: vi.fn().mockResolvedValue(undefined), + save: vi.fn().mockResolvedValue(true), positionAt: vi.fn().mockReturnValue({ line: 0, character: 0 }), }, }) @@ -1853,4 +1858,714 @@ describe("DiffViewProvider", () => { expect(vscode.window.showTextDocument).toHaveBeenCalled() }) }) + describe("removeAdoptedDirectories", () => { + it("removes the recorded directories deepest first and drops the tracking", async () => { + const dirs = [mockCwd + "/early", mockCwd + "/early/nested"] + diffViewProvider.adoptCreatedDirectories(dirs) + + await diffViewProvider.removeAdoptedDirectories() + + expect(vi.mocked(fs.rmdir).mock.calls.map((c) => c[0])).toEqual([...dirs].reverse()) + expect(diffViewProvider["createdDirs"]).toEqual([]) + }) + + it("does not reject when a directory cannot be removed and still attempts the rest", async () => { + const dirs = [mockCwd + "/early", mockCwd + "/early/nested"] + const failure = Object.assign(new Error("EPERM: operation not permitted"), { code: "EPERM" }) + vi.mocked(fs.rmdir).mockRejectedValueOnce(failure) + diffViewProvider.adoptCreatedDirectories(dirs) + + await expect(diffViewProvider.removeAdoptedDirectories()).resolves.toBeUndefined() + + expect(fs.rmdir).toHaveBeenCalledTimes(2) + expect(fs.rmdir).toHaveBeenCalledWith(dirs[0]) + }) + + it("reports a removal failure that is not an expected cleanup condition", async () => { + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + const dir = mockCwd + "/early/nested" + vi.mocked(fs.rmdir).mockRejectedValueOnce(Object.assign(new Error("EACCES"), { code: "EACCES" })) + diffViewProvider.adoptCreatedDirectories([dir]) + + await diffViewProvider.removeAdoptedDirectories() + + expect(errorSpy).toHaveBeenCalledWith( + "Error removing a directory this edit created:", + dir, + expect.any(Error), + ) + errorSpy.mockRestore() + }) + + it("stays quiet about a directory that is already gone or no longer empty", async () => { + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + vi.mocked(fs.rmdir) + .mockRejectedValueOnce(Object.assign(new Error("ENOENT"), { code: "ENOENT" })) + .mockRejectedValueOnce(Object.assign(new Error("ENOTEMPTY"), { code: "ENOTEMPTY" })) + diffViewProvider.adoptCreatedDirectories([mockCwd + "/early", mockCwd + "/early/nested"]) + + await diffViewProvider.removeAdoptedDirectories() + + expect(fs.rmdir).toHaveBeenCalledTimes(2) + expect(errorSpy).not.toHaveBeenCalled() + errorSpy.mockRestore() + }) + + it("does nothing when no directory was recorded", async () => { + await expect(diffViewProvider.removeAdoptedDirectories()).resolves.toBeUndefined() + + expect(fs.rmdir).not.toHaveBeenCalled() + }) + }) + + describe("discardUnapprovedStream method", () => { + const mockTargetPath = `${mockCwd}/mock-target-file.ts` + + const makeAbandonedDocument = (callOrder: string[], isDirty = true) => ({ + isDirty, + getText: () => "partial model output", + positionAt: (offset: number) => ({ line: 0, character: offset }), + uri: { fsPath: mockTargetPath, path: mockTargetPath }, + // TextDocument.save() resolves true on success; the discard reads that result. + save: vi.fn(async () => { + callOrder.push("save") + return true + }), + }) + + const openAbandonedView = (document: unknown, createdDirs: string[], callOrder: string[]) => + Object.assign(diffViewProvider, { + relPath: "mock-target-file.ts", + activeDiffEditor: { document }, + editType: "create", + createdDirs, + // open() records the placeholder it wrote; the discard may only delete what it owns. + placeholderPath: `${mockCwd}/mock-target-file.ts`, + closeAllDiffViews: vi.fn(async () => { + callOrder.push("closeDiffViews") + }), + closeFileTab: vi.fn(async () => { + callOrder.push("closeFileTab") + }), + }) + + it("never persists the unapproved buffer: blanks it with an empty replacement, then deletes the placeholder", async () => { + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder) + openAbandonedView(document, [], callOrder) + vi.mocked(vscode.workspace.applyEdit).mockImplementation(async () => { + callOrder.push("applyEdit") + return true + }) + + await diffViewProvider.discardUnapprovedStream() + + // The diff tabs are closed only after the buffer is clean: closeAllDiffViews() skips + // dirty tabs, so closing first left the vscode.diff tab open over the unlink below. + expect(callOrder).toEqual(["applyEdit", "save", "closeDiffViews", "closeFileTab"]) + // The only content the discard may write is an empty buffer, and it must be + // written to THIS document before the save that would otherwise persist the + // partial model output. + expect(mockWorkspaceEdit.replace).toHaveBeenCalledTimes(1) + const [replacedUri, , replacedText] = mockWorkspaceEdit.replace.mock.calls[0] + expect(replacedUri).toEqual(document.uri) + // The replacement covers the whole buffer: the range is built from + // positionAt(0) to positionAt(getText().length), so nothing can survive it. + expect(vi.mocked(vscode.Range)).toHaveBeenCalledWith( + { line: 0, character: 0 }, + { line: 0, character: "partial model output".length }, + ) + expect(replacedText).toBe("") + expect(mockWorkspaceEdit.delete).not.toHaveBeenCalled() + // The placeholder for THIS relPath is what gets removed. + expect(fs.unlink).toHaveBeenCalledWith(expect.stringContaining("mock-target-file.ts")) + }) + + it("restores a modify buffer in memory and never saves the target file", async () => { + // revertChanges() restores a modify by applyEdit + document.save(). For a stream the + // user never approved that save is a write they never asked for, and for a + // .rooignore-denied path a write the policy forbids outright, so the discard restores + // the buffer and leaves the file on disk untouched. + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder) + openAbandonedView(document, [], callOrder) + Object.assign(diffViewProvider, { + editType: "modify", + originalContent: "original content", + // A modify owns no placeholder: only open() of a create writes one. + placeholderPath: undefined, + }) + vi.mocked(vscode.workspace.applyEdit).mockImplementation(async () => { + callOrder.push("applyEdit") + return true + }) + + await diffViewProvider.discardUnapprovedStream() + + expect(callOrder).toEqual(["applyEdit", "closeDiffViews", "closeFileTab"]) + const [replacedUri, , replacedText] = mockWorkspaceEdit.replace.mock.calls[0] + expect(replacedUri).toEqual(document.uri) + expect(replacedText).toBe("original content") + // The replacement covers the whole buffer, so no streamed content survives it. + expect(vi.mocked(vscode.Range)).toHaveBeenCalledWith( + { line: 0, character: 0 }, + { line: 0, character: "partial model output".length }, + ) + expect(document.save).not.toHaveBeenCalled() + // Nothing is unlinked for a modify: saveChanges() may have written this very file with + // the user's approval, and this edit created no directories. + expect(fs.unlink).not.toHaveBeenCalled() + expect(fs.rmdir).not.toHaveBeenCalled() + }) + + it("keeps the placeholder on disk when the buffer restore fails to apply, and reports the hazard", async () => { + // applyEdit resolving false leaves the unapproved partial content in the buffer. Saving + // then would persist exactly what this method exists to discard, so the save is skipped + // and the caller is told the rollback did not happen. + // + // The placeholder is deliberately NOT removed here: the tab is still open and still + // dirty, so unlinking would leave it pointing at a path with no file behind it and the + // next Ctrl+S would recreate the file with the unapproved content. + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder) + openAbandonedView(document, [], callOrder) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(false) + + await expect(diffViewProvider.discardUnapprovedStream()).rejects.toThrow( + /could not restore the diff editor buffer/i, + ) + + expect(document.save).not.toHaveBeenCalled() + expect(fs.unlink).not.toHaveBeenCalled() + expect(fs.rmdir).not.toHaveBeenCalled() + }) + + it("keeps the placeholder when saving the restored buffer rejects, so the dirty tab keeps its file", async () => { + // The restore applied but the save of the emptied placeholder failed: the buffer is still + // dirty and its tab still open, which is the same hazard as an applyEdit that resolved + // false - unlinking would orphan the editor and the next Ctrl+S would resurrect the + // unapproved content into a file that no longer exists. + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder) + document.save.mockRejectedValueOnce(new Error("save rejected")) + openAbandonedView(document, [`${mockCwd}/mock-dir`], callOrder) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + + await expect(diffViewProvider.discardUnapprovedStream()).rejects.toThrow("save rejected") + + expect(fs.unlink).not.toHaveBeenCalled() + expect(fs.rmdir).not.toHaveBeenCalled() + }) + + it("removes the directories this edit created in reverse order, after the placeholder unlink", async () => { + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder) + const createdDirs = [`${mockCwd}/mock-dir`, `${mockCwd}/mock-dir/nested`] + openAbandonedView(document, createdDirs, callOrder) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + + await diffViewProvider.discardUnapprovedStream() + + expect(fs.unlink).toHaveBeenCalledTimes(1) + expect(fs.rmdir).toHaveBeenCalledTimes(createdDirs.length) + // Reverse order, so a parent is never removed before the child inside it, and + // only the directories this edit created - never a pre-existing one. + expect(vi.mocked(fs.rmdir).mock.calls.map((c) => c[0])).toEqual([...createdDirs].reverse()) + const unlinkOrder = vi.mocked(fs.unlink).mock.invocationCallOrder[0] + for (const order of vi.mocked(fs.rmdir).mock.invocationCallOrder) { + expect(order).toBeGreaterThan(unlinkOrder) + } + }) + + it("does not save when the abandoned buffer is already clean", async () => { + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder, false) + openAbandonedView(document, [], callOrder) + + await diffViewProvider.discardUnapprovedStream() + + expect(vscode.workspace.applyEdit).not.toHaveBeenCalled() + expect(document.save).not.toHaveBeenCalled() + expect(mockWorkspaceEdit.replace).not.toHaveBeenCalled() + expect(fs.unlink).toHaveBeenCalledTimes(1) + }) + + it("keeps the artifacts when the editor work rejects with a dirty buffer, and reports the failure", async () => { + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder) + const createdDirs = [`${mockCwd}/mock-dir`] + openAbandonedView(document, createdDirs, callOrder) + vi.mocked(vscode.workspace.applyEdit).mockRejectedValue(new Error("applyEdit rejected")) + + // The caller (discardUnapprovedStreamBeforeReset) treats a throw as "rollback hazard" + // and reports it, so the failure must surface. The artifacts stay behind on purpose: + // the buffer is still dirty and its tab still open, and removing the directory would + // orphan the file the open editor is pointing at. + await expect(diffViewProvider.discardUnapprovedStream()).rejects.toThrow("applyEdit rejected") + + expect(fs.unlink).not.toHaveBeenCalled() + expect(fs.rmdir).not.toHaveBeenCalled() + }) + + it("removes the artifacts when the editor work rejects after the buffer was restored clean", async () => { + // A clean buffer means no live editor can resurrect the content, so the placeholder + // and the directories really are debris and must still go even though the editor work + // failed - this is the case the original cleanup ordering existed for. + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder, false) + openAbandonedView(document, [`${mockCwd}/mock-dir`], callOrder) + Object.assign(diffViewProvider, { + closeFileTab: vi.fn(async () => { + throw new Error("close rejected") + }), + }) + + await expect(diffViewProvider.discardUnapprovedStream()).rejects.toThrow("close rejected") + + expect(fs.unlink).toHaveBeenCalledWith(expect.stringContaining("mock-target-file.ts")) + expect(fs.rmdir).toHaveBeenCalledWith(`${mockCwd}/mock-dir`) + }) + + it("removes the placeholder and created directories even when open() failed before activeDiffEditor was assigned", async () => { + // open() creates the directories (line 130) and the empty placeholder (line 134) + // BEFORE openDiffEditor() assigns activeDiffEditor (line 173). If that call rejects + // - the 10s timeout or a failed vscode.diff - an abandoned create therefore has + // artifacts on disk and no editor, and the early return must not skip the cleanup. + Object.assign(diffViewProvider, { + relPath: "mock-target-file.ts", + activeDiffEditor: undefined, + editType: "create", + createdDirs: [`${mockCwd}/mock-dir`], + // open() had already written the placeholder before openDiffEditor() rejected, + // so this edit owns the path. + placeholderPath: `${mockCwd}/mock-target-file.ts`, + }) + + await diffViewProvider.discardUnapprovedStream() + + expect(fs.unlink).toHaveBeenCalledWith(expect.stringContaining("mock-target-file.ts")) + expect(fs.rmdir).toHaveBeenCalledWith(`${mockCwd}/mock-dir`) + // Nothing editor-side ran, and nothing threw: the artifacts are the whole job here. + expect(vscode.workspace.applyEdit).not.toHaveBeenCalled() + // Bracket access for the private field: the snapshot must be consumed, so a second + // discard cannot try to remove the same directories again. + expect(diffViewProvider["createdDirs"]).toEqual([]) + }) + + it("logs and rejects when the placeholder unlink fails for a reason other than ENOENT", async () => { + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder) + openAbandonedView(document, [], callOrder) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + const permissionError = Object.assign(new Error("EPERM: operation not permitted"), { code: "EPERM" }) + vi.mocked(fs.unlink).mockRejectedValueOnce(permissionError) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + await expect(diffViewProvider.discardUnapprovedStream()).rejects.toThrow("EPERM: operation not permitted") + + expect(errorSpy).toHaveBeenCalledWith("Error removing abandoned write_to_file artifacts:", permissionError) + }) + + it("logs and rejects when removing a created directory fails for a reason other than ENOENT", async () => { + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder) + openAbandonedView(document, [`${mockCwd}/mock-dir`], callOrder) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + vi.mocked(fs.unlink).mockResolvedValue(undefined) + const notEmpty = Object.assign(new Error("ENOTEMPTY: directory not empty"), { code: "ENOTEMPTY" }) + vi.mocked(fs.rmdir).mockRejectedValueOnce(notEmpty) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + await expect(diffViewProvider.discardUnapprovedStream()).rejects.toThrow("ENOTEMPTY: directory not empty") + + expect(errorSpy).toHaveBeenCalledWith("Error removing abandoned write_to_file artifacts:", notEmpty) + }) + + it("tolerates an already-deleted placeholder and directories (ENOENT) without throwing", async () => { + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder) + // Two created directories: the tolerance has to hold for every one of them, so the + // rejection is keyed on the path rather than on which call comes first. + const dirA = `${mockCwd}/mock-dir-a` + const dirB = `${mockCwd}/nested/mock-dir-b` + openAbandonedView(document, [dirA, dirB], callOrder) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + const enoent = Object.assign(new Error("ENOENT"), { code: "ENOENT" }) + // A single unlink call in this flow, so the placeholder stays Once-based. + vi.mocked(fs.unlink).mockRejectedValueOnce(enoent) + vi.mocked(fs.rmdir).mockImplementation(async (dir) => { + if (dir === dirA || dir === dirB) { + throw enoent + } + }) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + await expect(diffViewProvider.discardUnapprovedStream()).resolves.toBeUndefined() + + // Both directories were attempted; the assertion is on the paths, not the order. + expect(fs.rmdir).toHaveBeenCalledWith(dirA) + expect(fs.rmdir).toHaveBeenCalledWith(dirB) + expect(errorSpy).not.toHaveBeenCalledWith( + "Error removing abandoned write_to_file artifacts:", + expect.anything(), + ) + + // beforeEach only clears call data, not implementations, so put the default back. + vi.mocked(fs.rmdir).mockResolvedValue(undefined) + }) + + it("reports the editor failure rather than the cleanup failure when both happen", async () => { + const callOrder: string[] = [] + // A clean buffer keeps the cleanup in play after a failed editor step, which is the + // only way both failures can coexist: an editor failure that leaves the buffer dirty + // deliberately skips the cleanup instead of racing it. + const document = makeAbandonedDocument(callOrder, false) + openAbandonedView(document, [`${mockCwd}/mock-dir`], callOrder) + Object.assign(diffViewProvider, { + closeFileTab: vi.fn(async () => { + throw new Error("close rejected") + }), + }) + vi.mocked(fs.unlink).mockRejectedValueOnce(Object.assign(new Error("EPERM"), { code: "EPERM" })) + vi.spyOn(console, "error").mockImplementation(() => {}) + + // The caller reports a rollback hazard from the thrown error, so the ORIGINAL failure + // is what must surface; the cleanup failure is still logged for the operator. + await expect(diffViewProvider.discardUnapprovedStream()).rejects.toThrow("close rejected") + }) + + it("treats a save the editor did not perform as a rollback failure and keeps the artifacts", async () => { + // TextDocument.save() resolves false when the editor did not write. Reading that as + // success deleted the placeholder under a still-dirty tab and still reported a + // restored preview - the next Ctrl+S recreates the file holding exactly the content + // this method exists to discard. + const document = { + isDirty: true, + getText: () => "partial model output", + positionAt: (offset: number) => ({ line: 0, character: offset }), + uri: { fsPath: mockTargetPath, path: mockTargetPath }, + save: vi.fn().mockResolvedValue(false), + } + openAbandonedView(document, [mockCwd + "/mock-dir"], []) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + + await expect(diffViewProvider.discardUnapprovedStream()).rejects.toThrow( + "Could not save the restored diff editor buffer", + ) + + expect(fs.unlink).not.toHaveBeenCalled() + expect(fs.rmdir).not.toHaveBeenCalled() + expect(document.save).toHaveBeenCalledTimes(1) + }) + + it("leaves a file this edit never created a placeholder for on disk", async () => { + // reset() does not clear relPath, and saveDirectly() sets it for an approved write. + // Without ownership tracking the discard would unlink that approved file: an approved + // write_to_file to A, then a new write_to_file whose block exits through a rollback + // path before open() ever ran, still sees relPath === A. + const callOrder: string[] = [] + const document = makeAbandonedDocument(callOrder) + Object.assign(diffViewProvider, { + relPath: "previously-approved-file.ts", + activeDiffEditor: { document }, + editType: undefined, + createdDirs: [], + placeholderPath: undefined, + }) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + + await expect(diffViewProvider.discardUnapprovedStream()).resolves.toBeUndefined() + + expect(fs.unlink).not.toHaveBeenCalled() + expect(fs.rmdir).not.toHaveBeenCalled() + // The buffer is still blanked and the tab closed: the unapproved content is gone + // without touching a file the user approved. + expect(mockWorkspaceEdit.replace).toHaveBeenCalledTimes(1) + expect(mockWorkspaceEdit.replace.mock.calls[0][2]).toBe("") + }) + + const ownedSaveSetup = (save: ReturnType) => { + const document = { + uri: { fsPath: mockTargetPath }, + getText: vi.fn().mockReturnValue("approved content"), + lineCount: 1, + positionAt: (offset: number) => ({ line: 0, character: offset }), + isDirty: true, + save, + } + Object.assign(diffViewProvider, { + relPath: "mock-target-file.ts", + newContent: "approved content", + editType: "create", + activeDiffEditor: { + document, + selection: { active: { line: 0, character: 0 }, anchor: { line: 0, character: 0 } }, + }, + createdDirs: [], + placeholderPath: mockTargetPath, + closeAllDiffViews: vi.fn().mockResolvedValue(undefined), + keepOrCloseEditedFile: vi.fn().mockResolvedValue(undefined), + restorePreviewTabs: vi.fn().mockResolvedValue(undefined), + }) + } + + it("keeps placeholder ownership when the approved save rejects, so the teardown can still clean up", async () => { + // saveChanges() is the approved write. Until it lands, the empty file open() created is + // still an artifact this edit owns: releasing ownership first told the discard that + // nothing was left to remove, and a rejected save left an empty or half-written new + // file on disk with nobody able to delete it. + ownedSaveSetup(vi.fn().mockRejectedValue(new Error("save rejected"))) + + await expect(diffViewProvider.saveChanges(false, 0)).rejects.toThrow("save rejected") + + expect(diffViewProvider["placeholderPath"]).toBe(mockTargetPath) + }) + + it("releases placeholder ownership once the approved save lands", async () => { + // The approved content replaced the placeholder, so the path is no longer this edit's + // to delete: a later rollback must not unlink content the user accepted. + const save = vi.fn().mockResolvedValue(undefined) + ownedSaveSetup(save) + + await expect(diffViewProvider.saveChanges(false, 0)).resolves.toBeDefined() + + expect(save).toHaveBeenCalledTimes(1) + expect(diffViewProvider["placeholderPath"]).toBeUndefined() + }) + + it("removes the placeholder that a new-file open() wrote, driven through the public methods", async () => { + // The ownership lifecycle lives at the public-method layer, so exercise it there: + // open() writes the empty placeholder and records it, and the discard of an abandoned + // create removes exactly that path. + const relPath = "owned-placeholder.ts" + const fsPath = `${mockCwd}/${relPath}` + const mockEditor = { + document: { + uri: { fsPath, scheme: "file" }, + getText: vi.fn().mockReturnValue(""), + isDirty: false, + save: vi.fn().mockResolvedValue(true), + lineCount: 0, + }, + selection: { active: { line: 0, character: 0 }, anchor: { line: 0, character: 0 } }, + edit: vi.fn().mockResolvedValue(true), + revealRange: vi.fn(), + } + // Structural double for the mocked editor: the mock only implements the members + // open() touches, so it is routed through unknown rather than any. + const editor = mockEditor as unknown as vscode.TextEditor + vi.mocked(vscode.window).visibleTextEditors = [editor] + vi.mocked(vscode.window.showTextDocument).mockResolvedValue(editor) + vi.mocked(vscode.workspace.onDidOpenTextDocument).mockImplementation((callback) => { + setTimeout(() => callback({ uri: { fsPath, scheme: "file" } } as vscode.TextDocument), 0) + return { dispose: vi.fn() } + }) + vi.mocked(vscode.window.onDidChangeVisibleTextEditors).mockReturnValue({ dispose: vi.fn() }) + vi.mocked(vscode.window.onDidChangeTextEditorVisibleRanges).mockReturnValue({ dispose: vi.fn() }) + vi.mocked(vscode.languages.getDiagnostics).mockReturnValue([]) + diffViewProvider.editType = "create" + + await diffViewProvider.open(relPath) + + // open() created the placeholder, and this edit now owns it. + expect(fs.writeFile).toHaveBeenCalledWith(fsPath, "") + + await diffViewProvider.discardUnapprovedStream() + + expect(fs.unlink).toHaveBeenCalledWith(fsPath) + }) + + it("keeps directories created before open() visible to the discard", async () => { + // handlePartial() creates the parent directories before the diff view exists and hands + // them over; open() then finds them already there and records none of its own. An + // assigning open() would overwrite the list and hand those directories back to the + // leak: the discard removes createdDirs and reset() just drops it. + const relPath = "early-directories.ts" + const fsPath = `${mockCwd}/${relPath}` + const early = [`${mockCwd}/early/nested`] + const mockEditor = { + document: { + uri: { fsPath, scheme: "file" }, + getText: vi.fn().mockReturnValue(""), + isDirty: false, + save: vi.fn().mockResolvedValue(true), + lineCount: 0, + }, + selection: { active: { line: 0, character: 0 }, anchor: { line: 0, character: 0 } }, + edit: vi.fn().mockResolvedValue(true), + revealRange: vi.fn(), + } + // Structural double for the mocked editor: the mock only implements the members + // open() touches, so it is routed through unknown rather than any. + const editor = mockEditor as unknown as vscode.TextEditor + vi.mocked(vscode.window).visibleTextEditors = [editor] + vi.mocked(vscode.window.showTextDocument).mockResolvedValue(editor) + vi.mocked(vscode.workspace.onDidOpenTextDocument).mockImplementation((callback) => { + setTimeout(() => callback({ uri: { fsPath, scheme: "file" } } as vscode.TextDocument), 0) + return { dispose: vi.fn() } + }) + vi.mocked(vscode.window.onDidChangeVisibleTextEditors).mockReturnValue({ dispose: vi.fn() }) + vi.mocked(vscode.window.onDidChangeTextEditorVisibleRanges).mockReturnValue({ dispose: vi.fn() }) + vi.mocked(vscode.languages.getDiagnostics).mockReturnValue([]) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + diffViewProvider.editType = "create" + diffViewProvider.adoptCreatedDirectories(early) + + await diffViewProvider.open(relPath) + + await diffViewProvider.discardUnapprovedStream() + + expect(fs.rmdir).toHaveBeenCalledWith(early[0]) + }) + + it("does not remove the file once saveChanges() has approved the content", async () => { + // saveChanges() turns the placeholder into approved content and drops this edit's + // claim on it, so a later abandoned cleanup must leave the approved file on disk. + const relPath = "approved-by-save.ts" + const fsPath = `${mockCwd}/${relPath}` + Object.assign(diffViewProvider, { + relPath, + newContent: "approved content", + editType: "create", + createdDirs: [], + placeholderPath: fsPath, + preDiagnostics: [], + activeDiffEditor: { + document: { + uri: { fsPath, scheme: "file" }, + getText: vi.fn().mockReturnValue("approved content"), + isDirty: false, + save: vi.fn().mockResolvedValue(true), + }, + }, + closeAllDiffViews: vi.fn().mockResolvedValue(undefined), + closeFileTab: vi.fn().mockResolvedValue(undefined), + }) + vi.mocked(vscode.languages.getDiagnostics).mockReturnValue([]) + vi.mocked(vscode.window.showTextDocument).mockResolvedValue(undefined as never) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + + await diffViewProvider.saveChanges(false, 0) + + await diffViewProvider.discardUnapprovedStream() + + expect(fs.unlink).not.toHaveBeenCalled() + }) + + it("does not claim ownership of a placeholder whose write failed", async () => { + // open() writes the empty placeholder and only then records it as this edit's. When the + // write itself fails, this edit created nothing, so a later discard must not unlink a + // file it never made - claiming ownership before the write would. + const relPath = "unwritten-placeholder.ts" + const fsPath = `${mockCwd}/${relPath}` + const callOrder: string[] = [] + vi.mocked(fs.writeFile).mockRejectedValueOnce(new Error("EDQUOT: quota exceeded")) + diffViewProvider.editType = "create" + + await expect(diffViewProvider.open(relPath)).rejects.toThrow("EDQUOT") + + expect(diffViewProvider["placeholderPath"]).toBeUndefined() + + // A discard that runs afterwards for the same relPath has no ownership to act on. + Object.assign(diffViewProvider, { + relPath, + activeDiffEditor: { document: makeAbandonedDocument(callOrder) }, + createdDirs: [], + closeAllDiffViews: vi.fn().mockResolvedValue(undefined), + closeFileTab: vi.fn().mockResolvedValue(undefined), + }) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + + await diffViewProvider.discardUnapprovedStream() + + expect(fs.unlink).not.toHaveBeenCalled() + }) + + it("releases placeholder ownership once revertChanges() has deleted the new file", async () => { + // revertChanges() unlinks the file a create made. Ownership ends with that delete: a + // later discard that still claimed the path would unlink a file the next edit may + // already have recreated. + const relPath = "reverted-placeholder.ts" + const fsPath = `${mockCwd}/${relPath}` + Object.assign(diffViewProvider, { + relPath, + editType: "create", + createdDirs: [], + placeholderPath: fsPath, + activeDiffEditor: { + document: { + uri: { fsPath, scheme: "file" }, + getText: vi.fn().mockReturnValue("partial content"), + isDirty: false, + save: vi.fn().mockResolvedValue(true), + }, + }, + closeAllDiffViews: vi.fn().mockResolvedValue(undefined), + closeFileTab: vi.fn().mockResolvedValue(undefined), + }) + + await diffViewProvider.revertChanges() + + expect(fs.unlink).toHaveBeenCalledWith(fsPath) + expect(diffViewProvider["placeholderPath"]).toBeUndefined() + + // A discard after a successful revert must not reach for that path a second time. + const callOrder: string[] = [] + Object.assign(diffViewProvider, { activeDiffEditor: { document: makeAbandonedDocument(callOrder) } }) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + + await diffViewProvider.discardUnapprovedStream() + + expect(fs.unlink).toHaveBeenCalledTimes(1) + }) + + it("releases the placeholder it already deleted when the rest of the revert fails", async () => { + // revertChanges() ends in reset(), which would drop the ownership claim anyway. When a + // later step of the revert throws, reset() never runs: the file this revert already + // deleted must not stay claimed, or a discard after the failure unlinks whatever the + // next edit recreated at that path. + const relPath = "reverted-then-rmdir-failed.ts" + const fsPath = `${mockCwd}/${relPath}` + Object.assign(diffViewProvider, { + relPath, + editType: "create", + createdDirs: [`${mockCwd}/leftover-dir`], + placeholderPath: fsPath, + activeDiffEditor: { + document: { + uri: { fsPath, scheme: "file" }, + getText: vi.fn().mockReturnValue("partial content"), + isDirty: false, + save: vi.fn().mockResolvedValue(true), + }, + }, + closeAllDiffViews: vi.fn().mockResolvedValue(undefined), + closeFileTab: vi.fn().mockResolvedValue(undefined), + }) + vi.mocked(fs.rmdir).mockRejectedValueOnce(new Error("ENOTEMPTY: directory not empty")) + + await expect(diffViewProvider.revertChanges()).rejects.toThrow("ENOTEMPTY") + + expect(fs.unlink).toHaveBeenCalledWith(fsPath) + // The delete landed, so the claim ended with it - reset() never ran here. + expect(diffViewProvider["placeholderPath"]).toBeUndefined() + + // A discard after the failed revert must not reach for that path a second time. + const callOrder: string[] = [] + Object.assign(diffViewProvider, { activeDiffEditor: { document: makeAbandonedDocument(callOrder) } }) + vi.mocked(vscode.workspace.applyEdit).mockResolvedValue(true) + + await diffViewProvider.discardUnapprovedStream() + + expect(fs.unlink).toHaveBeenCalledTimes(1) + }) + + it("does nothing when no abandoned view is open", async () => { + Object.assign(diffViewProvider, { relPath: undefined, activeDiffEditor: undefined }) + + await diffViewProvider.discardUnapprovedStream() + + expect(fs.unlink).not.toHaveBeenCalled() + }) + }) })