[Fix] Subtasks fail to return when users work across windows

packages/core/package.json

@@ -5,6 +5,7 @@
 	"type": "module",
 	"exports": {
 		".": "./src/index.ts",
+		"./async-task-tracker": "./src/async-task-tracker.ts",
 		"./cli": "./src/cli.ts",
 		"./browser": "./src/browser.ts"
 	},

packages/core/src/async-task-tracker.spec.ts

@@ -0,0 +1,65 @@
+import { AsyncTaskTracker } from "./async-task-tracker.js"
+
+describe("AsyncTaskTracker", () => {
+	it("drains resolving and rejecting tasks", async () => {
+		const tracker = new AsyncTaskTracker()
+		let resolveTask!: () => void
+		let rejectTask!: (error: Error) => void
+		const resolving = new Promise<void>((resolve) => (resolveTask = resolve))
+		const rejecting = new Promise<void>((_resolve, reject) => (rejectTask = reject))
+
+		expect(tracker.track(resolving)).toBe(resolving)
+		expect(tracker.track(rejecting)).toBe(rejecting)
+		const drained = vi.fn()
+		const drain = tracker.drain().then(drained)
+		await Promise.resolve()
+		expect(drained).not.toHaveBeenCalled()
+
+		resolveTask()
+		rejectTask(new Error("expected rejection"))
+		await drain
+		expect(drained).toHaveBeenCalledOnce()
+	})
+
+	it("includes tasks tracked while a drain is in progress", async () => {
+		const tracker = new AsyncTaskTracker()
+		let resolveFirst!: () => void
+		let resolveSecond!: () => void
+		tracker.track(new Promise<void>((resolve) => (resolveFirst = resolve)))
+		const drain = tracker.drain()
+		tracker.track(new Promise<void>((resolve) => (resolveSecond = resolve)))
+
+		resolveFirst()
+		let drained = false
+		void drain.then(() => (drained = true))
+		await Promise.resolve()
+		expect(drained).toBe(false)
+		resolveSecond()
+		await drain
+		expect(drained).toBe(true)
+	})
+
+	it("stops guarded callbacks before draining tracked work", async () => {
+		const tracker = new AsyncTaskTracker()
+		const callback = vi.fn(async (value: string) => value.length)
+		let resolveTask!: () => void
+		const task = tracker.track(new Promise<void>((resolve) => (resolveTask = resolve)))
+		expect(tracker.isActive).toBe(true)
+		await expect(tracker.runIfActive(callback, "active")).resolves.toBe(6)
+		await expect(tracker.trackIfActive(async () => "tracked")).resolves.toBe("tracked")
+
+		let drained = false
+		const close = tracker.closeAndDrain().then(() => (drained = true))
+		await Promise.resolve()
+		expect(drained).toBe(false)
+		expect(tracker.isActive).toBe(false)
+		await expect(tracker.runIfActive(callback, "closed")).resolves.toBeUndefined()
+		await expect(tracker.trackIfActive(async () => "closed")).rejects.toThrow("tracker is closed")
+		expect(callback).toHaveBeenCalledOnce()
+
+		resolveTask()
+		await task
+		await close
+		expect(drained).toBe(true)
+	})
+})

packages/core/src/async-task-tracker.ts

@@ -0,0 +1,35 @@
+export class AsyncTaskTracker {
+	private readonly tasks = new Set<Promise<unknown>>()
+	private active = true
+
+	get isActive(): boolean {
+		return this.active
+	}
+
+	track<T>(task: Promise<T>): Promise<T> {
+		this.tasks.add(task)
+		void task.then(
+			() => this.tasks.delete(task),
+			() => this.tasks.delete(task),
+		)
+		return task
+	}
+
+	trackIfActive<T>(callback: () => Promise<T>): Promise<T> {
+		if (!this.active) return Promise.reject(new Error("Async task tracker is closed"))
+		return this.track(callback())
+	}
+
+	runIfActive<T, Result>(callback: (value: T) => Promise<Result>, value: T): Promise<Result | undefined> {
+		return this.active ? callback(value) : Promise.resolve(undefined)
+	}
+
+	async drain(): Promise<void> {
+		while (this.tasks.size) await Promise.allSettled(this.tasks)
+	}
+
+	async closeAndDrain(): Promise<void> {
+		this.active = false
+		await this.drain()
+	}
+}

packages/core/src/task-history/__tests__/task-history.spec.ts

@@ -2,7 +2,17 @@ import * as fs from "fs/promises"
 import * as os from "os"
 import * as path from "path"
 
-import { readTaskSessionsFromStoragePath } from "../index.js"
+import type { HistoryItem } from "@roo-code/types"
+
+import {
+	ABSENT_TASK_FILE_PREIMAGE,
+	INVALID_TASK_FILE_PREIMAGE,
+	expectedTaskFileStates,
+	isValidTaskFilePreImage,
+	matchesExpectedHistoryItem,
+	readTaskSessionsFromStoragePath,
+	taskFilePreImage,
+} from "../index.js"
 
 describe("readTaskSessionsFromStoragePath", () => {
 	let tempDir: string
@@ -112,3 +122,37 @@ describe("readTaskSessionsFromStoragePath", () => {
 		await expect(readTaskSessionsFromStoragePath(tempDir)).resolves.toEqual([])
 	})
 })
+
+describe("task file pre-images", () => {
+	const item: HistoryItem = {
+		id: "task-1",
+		number: 1,
+		ts: 1,
+		task: "Task",
+		tokensIn: 0,
+		tokensOut: 0,
+		totalCost: 0,
+	}
+
+	it("distinguishes validated, absent, and invalid pre-images", () => {
+		const valid = taskFilePreImage(item, item.id, true)
+		expect(valid).toEqual(item)
+		expect(valid).not.toBe(item)
+		expect(taskFilePreImage(null, item.id, false)).toBe(ABSENT_TASK_FILE_PREIMAGE)
+		expect(taskFilePreImage(null, item.id, true)).toBe(INVALID_TASK_FILE_PREIMAGE)
+		expect(taskFilePreImage({ ...item, id: "other" }, item.id, true)).toBe(INVALID_TASK_FILE_PREIMAGE)
+		expect(taskFilePreImage({ ...item, id: "other" }, item.id, false)).toBe(INVALID_TASK_FILE_PREIMAGE)
+		expect(taskFilePreImage({ id: item.id }, item.id, true)).toBe(INVALID_TASK_FILE_PREIMAGE)
+	})
+
+	it("recognizes valid snapshots and expected records", () => {
+		expect(isValidTaskFilePreImage(item)).toBe(true)
+		expect(isValidTaskFilePreImage(ABSENT_TASK_FILE_PREIMAGE)).toBe(false)
+		expect(isValidTaskFilePreImage(INVALID_TASK_FILE_PREIMAGE)).toBe(false)
+		expect(expectedTaskFileStates(item, { ...item })).toEqual([item, item])
+		expect(expectedTaskFileStates(item, ABSENT_TASK_FILE_PREIMAGE)).toEqual([item])
+		expect(expectedTaskFileStates(item, INVALID_TASK_FILE_PREIMAGE)).toEqual([item])
+		expect(matchesExpectedHistoryItem(item, [{ ...item }], (left, right) => left.id === right.id)).toBe(true)
+		expect(matchesExpectedHistoryItem(item, [], (left, right) => left.id === right.id)).toBe(false)
+	})
+})

packages/core/src/task-history/index.ts

@@ -1,7 +1,7 @@
 import * as fs from "fs/promises"
 import * as path from "path"
 
-import type { HistoryItem } from "@roo-code/types"
+import { historyItemSchema, type HistoryItem } from "@roo-code/types"
 
 const HISTORY_ITEM_FILENAME = "history_item.json"
 const HISTORY_INDEX_FILENAME = "_index.json"
@@ -15,6 +15,34 @@ export interface TaskSessionEntry {
 	status?: HistoryItem["status"]
 }
 
+export const ABSENT_TASK_FILE_PREIMAGE = "absent" as const
+export const INVALID_TASK_FILE_PREIMAGE = "invalid" as const
+export type TaskFilePreImage = HistoryItem | typeof ABSENT_TASK_FILE_PREIMAGE | typeof INVALID_TASK_FILE_PREIMAGE
+
+export function taskFilePreImage(existing: unknown, taskId: string, fileExists: boolean): TaskFilePreImage {
+	const parsed = historyItemSchema.safeParse(existing)
+	if (parsed.success && parsed.data.id === taskId) return structuredClone(existing as HistoryItem)
+	if (existing !== null) return INVALID_TASK_FILE_PREIMAGE
+	if (fileExists) return INVALID_TASK_FILE_PREIMAGE
+	return ABSENT_TASK_FILE_PREIMAGE
+}
+
+export function isValidTaskFilePreImage(preImage: TaskFilePreImage): preImage is HistoryItem {
+	return preImage !== ABSENT_TASK_FILE_PREIMAGE && preImage !== INVALID_TASK_FILE_PREIMAGE
+}
+
+export function expectedTaskFileStates(written: HistoryItem, preImage: TaskFilePreImage): HistoryItem[] {
+	return isValidTaskFilePreImage(preImage) ? [written, preImage] : [written]
+}
+
+export function matchesExpectedHistoryItem(
+	item: HistoryItem,
+	expected: readonly HistoryItem[],
+	equals: (left: HistoryItem, right: HistoryItem) => boolean,
+): boolean {
+	return expected.some((candidate) => equals(item, candidate))
+}
+
 function isRecord(value: unknown): value is Record<string, unknown> {
 	return typeof value === "object" && value !== null
 }