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/tools/WriteToFileTool.ts b/src/core/tools/WriteToFileTool.ts index ae026b4b86..c9776553b7 100644 --- a/src/core/tools/WriteToFileTool.ts +++ b/src/core/tools/WriteToFileTool.ts @@ -2,7 +2,7 @@ import path from "path" import delay from "delay" 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" @@ -23,9 +23,165 @@ 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) + } + + 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) + } + + private async resetDiffViewAfterWrite(task: Task): Promise { + await task.diffViewProvider.reset().catch((resetError) => { + console.error("Error resetting write_to_file diff view:", resetError) + }) + } + + /** + * Restore the diff editor document to its pre-streaming state and close the view. + * + * 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 revertChanges() relies on. No-op when no diff view + * is open. Failures are logged and swallowed so the remaining cleanup (reset, + * per-task state teardown) always continues. + */ + private async revertDiffChangesBeforeReset(task: Task): Promise { + await task.diffViewProvider.revertChanges().catch((revertError) => { + console.error("Error reverting write_to_file diff view changes:", revertError) + }) + } + + 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) + }) + } + + /** + * Teardown boundary for the handle() parse-failure path, where execute() never + * runs and therefore its finally (resetTaskPartialState) never runs either. + * + * Tears down the per-task 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 suppress the incidental + * parse error. + */ + override resetPartialState(): void { + super.resetPartialState() + for (const state of this.taskPartialStreamState.values()) { + state.task.off(RooCodeEventName.TaskAborted, state.abortCleanup) + } + this.taskPartialStreamState.clear() + } + async execute(params: WriteToFileParams, task: Task, callbacks: ToolCallbacks): Promise { const { pushToolResult, handleError, askApproval } = callbacks const relPath = params.path @@ -197,8 +353,13 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { const relPath: string | undefined = block.params.path const newContent: string | undefined = block.params.content + + // 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) + // Wait for path to stabilize before showing UI (prevents truncated paths) - if (!this.hasPathStabilized(relPath) || newContent === undefined) { + if (!this.hasPathStabilizedForTask(partialStreamState, relPath) || newContent === undefined) { return } 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..4506742c26 --- /dev/null +++ b/src/core/tools/__tests__/writeToFileTool-partial-state-cleanup.spec.ts @@ -0,0 +1,100 @@ +// 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> + } + finalizePartialToolAsk: 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), + }, + finalizePartialToolAsk: 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 reverting the diff document fails", async () => { + const task = buildTask("revert-fails", "inst-4") + const t = task as unknown as CleanupTask + t.diffViewProvider.revertChanges = vi.fn().mockRejectedValue(new Error("revert failed")) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + await writeToFileTool["revertDiffChangesBeforeReset"](task) + + expect(errorSpy).toHaveBeenCalledWith("Error reverting write_to_file diff view changes:", expect.any(Error)) + }) + + 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..bec63332a7 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" @@ -118,6 +119,9 @@ describe("writeToFileTool", () => { 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 +132,8 @@ describe("writeToFileTool", () => { return content }) + mockCline.taskId = "task-1" + mockCline.instanceId = "instance-1" mockCline.cwd = "/" mockCline.consecutiveMistakeCount = 0 mockCline.didEditFile = false @@ -186,8 +192,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 +429,111 @@ 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) + expect(mockCline.once).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, expect.any(Function)) + + 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 the base partial path and detaches every task's abort listener", async () => { + 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") + + // The base-class singleton field is reset by super.resetPartialState(). + writeToFileTool["lastSeenPartialPath"] = "stale-path" + writeToFileTool.resetPartialState() + + expect(writeToFileTool["lastSeenPartialPath"]).toBeUndefined() + expect(mockCline.off).toHaveBeenCalledWith(RooCodeEventName.TaskAborted, abortCleanup) + + // The per-task map was cleared too: a fresh delta sequence starts un-stabilized, so no + // second partial ask is issued. + await executeWriteFileTool({}, { isPartial: true }) + expect(mockCline.ask).toHaveBeenCalledTimes(1) + }) }) describe("user interaction", () => { @@ -470,6 +585,7 @@ describe("writeToFileTool", () => { // Second call with same path - path is now stabilized, error occurs await executeWriteFileTool({}, { isPartial: true }) expect(mockHandleError).toHaveBeenCalledWith("handling partial write_to_file", expect.any(Error)) + }) }) }) diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 874e2f8f08..75612323f9 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, @@ -636,6 +637,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) {