Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions src/core/task/Task.ts
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ import { buildNativeToolsArrayWithRestrictions } from "./build-tools"
import { ToolRepetitionDetector } from "../tools/ToolRepetitionDetector"
import { restoreTodoListForTask } from "../tools/UpdateTodoListTool"
import { FileContextTracker } from "../context-tracking/FileContextTracker"
import { ObservationRegistry } from "./observationRegistry"
import { RooIgnoreController } from "../ignore/RooIgnoreController"
import { RooProtectedController } from "../protect/RooProtectedController"
import { type AssistantMessageContent, presentAssistantMessage } from "../assistant-message"
Expand Down Expand Up @@ -296,6 +297,10 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
readonly instanceId: string
readonly metadata: TaskMetadata

// The observed on-disk version of each file this task has read. Declared here so the
// read tools can record it; a write guard later compares a token against this registry.
readonly observationRegistry = new ObservationRegistry()

todoList?: TodoItem[]

readonly rootTask: Task | undefined = undefined
Expand Down Expand Up @@ -392,6 +397,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {

providerRef: WeakRef<ClineProvider>
private readonly globalStoragePath: string
// Monotonic cancellation counter. `requestAbort()` and `requestDispose()` advance
// this fence, and the guarded-write path compares the value it captured when the
// write was queued with the value at publish time, so a queued write cannot publish
// after the task was cancelled or disposed.
cancellationGeneration: number = 0
currentRequestAbortController?: AbortController
/**
* Controller for the waiter on an in-flight `ensureModelFetched()` call (see
Expand Down Expand Up @@ -3595,6 +3605,13 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
if (isAbandoned) {
this.abandoned = true
}
// The S4a write guard compares this counter with the value a queued write
// captured when it started, so every cancellation must advance it and it
// must never go backwards. Main moved the abort flag into the run-state
// reducer (see the abort(this.runState) getter above), so this counter is
// the only thing left that separates a cancelled run's queued writes from
// the resumed run's, and only if it advances monotonically.
this.cancellationGeneration++
this.driveRunEvent({ tag: "abortRequested", abandoned: isAbandoned })
}

Expand All @@ -3604,6 +3621,10 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
}

requestDispose(): void {
// A host teardown never calls abortTask(): disposal has to raise the same
// counter, or a write still parked on its path chain publishes after the
// task is gone.
this.cancellationGeneration++
this.driveRunEvent({ tag: "disposeRequested" })
}

Expand Down Expand Up @@ -3632,6 +3653,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {

public abortTask(isAbandoned = false): Promise<void> {
this.requestAbort(isAbandoned)

this.cancelAssistantMessagePersistence()
this.abortPromise ??= this.abortTaskOnce()
return this.abortPromise
Expand Down Expand Up @@ -3754,6 +3776,13 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
console.error("Error removing event listeners:", error)
}

// A disposed task is no longer authoritative for what it read. The registry holds
// version tokens captured while the task was alive, and a disposed task can still be
// reachable through a parent/subtask reference; a guarded write must not accept one of
// those tokens for a file this task has not re-read since. Clearing also stops a long
// task from pinning every file it ever read.
this.observationRegistry.close()

// Release any terminals associated with this task.
try {
// Release any terminals associated with this task.
Expand Down
50 changes: 50 additions & 0 deletions src/core/task/__tests__/Task.dispose.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import path from "node:path"
import { type ProviderSettings, RooCodeEventName } from "@roo-code/types"

import { Task } from "../Task"
import { ObservationRegistry } from "../observationRegistry"
import { ClineProvider } from "../../webview/ClineProvider"
import { OutputInterceptor } from "../../../integrations/terminal/OutputInterceptor"
import { providerIdentifiers } from "@roo-code/types/provider-identifiers"
Expand Down Expand Up @@ -118,6 +119,55 @@ describe("Task dispose method", () => {
expect(disposalComplete).toBe(true)
})

test("owns a per-Task observation registry", () => {
// ReadFileTool and the guarded-write path reach the registry only through the Task,
// so a Task that never built one - or shared one across Tasks - would silently drop
// read observations. The registry's own tests only exercise standalone instances, and
// the tool tests inject their own doubles, so nothing else pins this wiring.
expect(task.observationRegistry).toBeInstanceOf(ObservationRegistry)

const other = new Task({
provider: mockProvider as unknown as ClineProvider,
apiConfiguration: mockApiConfiguration,
startTask: false,
})
try {
expect(other.observationRegistry).toBeInstanceOf(ObservationRegistry)
expect(other.observationRegistry).not.toBe(task.observationRegistry)

// Observing in one Task must not be visible from another: parent and subtask
// authority over a file has to stay separate.
other.observationRegistry.observe("/workspace/a.ts", "v1", true)
expect(other.observationRegistry.get("/workspace/a.ts")?.version).toBe("v1")
expect(task.observationRegistry.has("/workspace/a.ts")).toBe(false)
expect(task.observationRegistry.size).toBe(0)
} finally {
void other.dispose().catch(() => {})
}
})

test("clears the per-Task observation registry on disposal", async () => {
// The registry holds on-disk version tokens for files this task read. Disposal does not
// free the Task object - a parent or subtask reference can outlive it - so a token
// captured before teardown must not survive as authority for a guarded write.
task.observationRegistry.observe("/workspace/a.ts", "v1", true)
task.observationRegistry.observe("/workspace/b.ts", "v2", false)
expect(task.observationRegistry.size).toBe(2)

await task.dispose()

expect(task.observationRegistry.size).toBe(0)
expect(task.observationRegistry.has("/workspace/a.ts")).toBe(false)
expect(task.observationRegistry.get("/workspace/b.ts")).toBeUndefined()

// A read that was still awaiting I/O when disposal began resumes afterwards. Its
// observation must not land in a retired registry, or the disposed Task would be
// authoritative for that file again.
expect(task.observationRegistry.closed).toBe(true)
task.observationRegistry.observe("/workspace/late.ts", "v3", true)
expect(task.observationRegistry.size).toBe(0)
})

test("should reject the memoized completion promise when disposal cannot start", async () => {
const disposalError = new Error("disposal failed")
skipCleanup = true
Expand Down
54 changes: 54 additions & 0 deletions src/core/task/__tests__/Task.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4769,6 +4769,57 @@ describe("Cline", () => {
).toHaveLength(1)
})

it("advances the cancellation generation on each cancellation and keeps it advanced across a resume", async () => {
// The S4a write guard compares the cancellation generation a queued write
// captured against the generation at publish time. The abort FLAG cannot do
// that on its own: it is set on cancellation and no resume path clears it,
// so on a resumed task it says only that some cancellation happened, not
// which run a queued write belongs to. Only the generation keeps the two
// runs apart, so it must never go backwards.
const task = new Task({
provider: mockProvider,
apiConfiguration: mockApiConfig,
task: "test task",
startTask: false,
})
vi.spyOn(task, "dispose").mockResolvedValue(undefined)

expect(task.cancellationGeneration).toBe(0)

await task.abortTask()
expect(task.abort).toBe(true)
expect(task.cancellationGeneration).toBe(1)

await task.resumeAfterDelegation()
// The run-state model (RSK-06) no longer clears the legacy `abort` flag
// on resume: that is exactly why the generation exists. Only the
// generation keeps the cancelled run and the resumed run apart.
expect(task.cancellationGeneration).toBe(1)

await task.abortTask()
expect(task.abort).toBe(true)
expect(task.cancellationGeneration).toBe(2)
})

it("advances the cancellation generation on disposal, without any explicit cancel", async () => {
// A task torn down by the host never sees abortTask(). Disposal has to raise the
// same generation, or a write still parked on its path chain publishes after the
// task is gone.
const task = new Task({
provider: mockProvider,
apiConfiguration: mockApiConfig,
task: "test task",
startTask: false,
})

expect(task.cancellationGeneration).toBe(0)

await task.dispose()

expect(task.abort).toBe(true)
expect(task.cancellationGeneration).toBe(1)
})

it("flushes pending state before TaskAborted and disposal while queue state is intact", async () => {
const task = new Task({
provider: mockProvider,
Expand Down Expand Up @@ -7762,6 +7813,9 @@ describe("Cline", () => {
// task directory to exist (uuid v7 is mocked to the fixed id below).
const taskDir = path.join(os.tmpdir(), "test-storage", "tasks", "00000000-0000-7000-8000-000000000000")
fsReal.mkdirSync(taskDir, { recursive: true })
// Earlier tests in this file can leave a ui_messages.json for the fixed task id;
// merge=true would read it here and fail the first save stage for an unrelated reason.
fsReal.rmSync(path.join(taskDir, "ui_messages.json"), { force: true })
// Declared outside the try so the finally block can restore them.
const updateSpy = vi
.spyOn(getTaskTestAccess(Task.prototype), "updateClineMessage")
Expand Down
132 changes: 132 additions & 0 deletions src/core/task/__tests__/observationRegistry.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
import { describe, it, expect, vi } from "vitest"

import { ObservationRegistry } from "../observationRegistry"

describe("ObservationRegistry", () => {
it("observe → get returns the recorded version and observedAt", () => {
const reg = new ObservationRegistry()
reg.observe("/a/b/c.ts", "1:2:300:4000000000:5000000000")

const obs = reg.get("/a/b/c.ts")
expect(obs).toBeDefined()
expect(obs!.version).toBe("1:2:300:4000000000:5000000000")
expect(typeof obs!.observedAt).toBe("number")
})

it("re-observe replaces the entry with a fresh observedAt", () => {
vi.useFakeTimers()
const reg = new ObservationRegistry()
reg.observe("/a/b/c.ts", "v1")
const first = reg.get("/a/b/c.ts")!
expect(first.version).toBe("v1")

vi.advanceTimersByTime(50)
reg.observe("/a/b/c.ts", "v2")
const second = reg.get("/a/b/c.ts")!
expect(second.version).toBe("v2")
expect(second.observedAt).toBeGreaterThan(first.observedAt)

vi.useRealTimers()
})

it("has returns true for observed paths, false otherwise", () => {
const reg = new ObservationRegistry()
reg.observe("/x.ts", "t1")
expect(reg.has("/x.ts")).toBe(true)
expect(reg.has("/y.ts")).toBe(false)
})

it("size reflects the number of observed entries", () => {
const reg = new ObservationRegistry()
expect(reg.size).toBe(0)
reg.observe("/a.ts", "t1")
reg.observe("/b.ts", "t2")
expect(reg.size).toBe(2)
})

it("clear removes all entries and resets size to 0", () => {
const reg = new ObservationRegistry()
reg.observe("/a.ts", "t1")
reg.observe("/b.ts", "t2")
reg.clear()
expect(reg.size).toBe(0)
expect(reg.get("/a.ts")).toBeUndefined()
expect(reg.has("/b.ts")).toBe(false)
})

it("get on empty registry returns undefined", () => {
const reg = new ObservationRegistry()
expect(reg.get("/any.ts")).toBeUndefined()
})

it("separate instances are independent — observing in one does not appear in the other", () => {
const regA = new ObservationRegistry()
const regB = new ObservationRegistry()
regA.observe("/shared.ts", "v1")
expect(regA.get("/shared.ts")).toBeDefined()
expect(regB.get("/shared.ts")).toBeUndefined()
regB.observe("/shared.ts", "v2")
expect(regA.get("/shared.ts")!.version).toBe("v1")
expect(regB.get("/shared.ts")!.version).toBe("v2")
})

describe("completeness scope (S4b follow-up #46)", () => {
it("defaults to a complete observation when the read scope is not given", () => {
const reg = new ObservationRegistry()
reg.observe("/a/b/c.ts", "v1")

expect(reg.get("/a/b/c.ts")!.complete).toBe(true)
})

it("records a partial observation when the read only returned a view of the file", () => {
const reg = new ObservationRegistry()
reg.observe("/a/b/c.ts", "v1", false)

expect(reg.get("/a/b/c.ts")!.complete).toBe(false)
})

it("re-observing replaces the entry's completeness with the new read's scope", () => {
const reg = new ObservationRegistry()
reg.observe("/a/b/c.ts", "v1", false)
reg.observe("/a/b/c.ts", "v2")

const obs = reg.get("/a/b/c.ts")!
expect(obs.version).toBe("v2")
expect(obs.complete).toBe(true)
})

it("re-observing with a partial scope downgrades a previously complete entry", () => {
const reg = new ObservationRegistry()
reg.observe("/a/b/c.ts", "v1")
reg.observe("/a/b/c.ts", "v2", false)

const obs = reg.get("/a/b/c.ts")!
expect(obs.version).toBe("v2")
expect(obs.complete).toBe(false)
})
})

describe("closure on task disposal", () => {
it("close drops every entry and reports the registry closed", () => {
const reg = new ObservationRegistry()
reg.observe("/a/b/c.ts", "v1")
expect(reg.closed).toBe(false)

reg.close()

expect(reg.size).toBe(0)
expect(reg.get("/a/b/c.ts")).toBeUndefined()
expect(reg.closed).toBe(true)
})

it("observe is a no-op after close, so a read resuming after disposal records nothing", () => {
const reg = new ObservationRegistry()
reg.close()

reg.observe("/late.ts", "v1")

expect(reg.size).toBe(0)
expect(reg.has("/late.ts")).toBe(false)
})
})
})
Loading
Loading