Add Cline events and roofactor subtasks slightly

e2e/src/suite/extension.test.ts

@@ -9,10 +9,6 @@ suite("Roo Code Extension", () => {
 	})
 
 	test("Commands should be registered", async () => {
-		const timeout = 10 * 1_000
-		const interval = 1_000
-		const startTime = Date.now()
-
 		const expectedCommands = [
 			"roo-cline.plusButtonClicked",
 			"roo-cline.mcpButtonClicked",
@@ -25,23 +21,6 @@ suite("Roo Code Extension", () => {
 			"roo-cline.improveCode",
 		]
 
-		while (Date.now() - startTime < timeout) {
-			const commands = await vscode.commands.getCommands(true)
-			const missingCommands = []
-
-			for (const cmd of expectedCommands) {
-				if (!commands.includes(cmd)) {
-					missingCommands.push(cmd)
-				}
-			}
-
-			if (missingCommands.length === 0) {
-				break
-			}
-
-			await new Promise((resolve) => setTimeout(resolve, interval))
-		}
-
 		const commands = await vscode.commands.getCommands(true)
 
 		for (const cmd of expectedCommands) {

e2e/src/suite/index.ts

@@ -8,55 +8,39 @@ import { RooCodeAPI } from "../../../src/exports/roo-code"
 import { waitUntilReady } from "./utils"
 
 declare global {
-	var extension: vscode.Extension<RooCodeAPI> | undefined
 	var api: RooCodeAPI
 }
 
 export async function run() {
-	const mocha = new Mocha({ ui: "tdd", timeout: 300_000 })
-	const testsRoot = path.resolve(__dirname, "..")
-
-	try {
-		// Find all test files.
-		const files = await glob("**/**.test.js", { cwd: testsRoot })
-		files.forEach((f: string) => mocha.addFile(path.resolve(testsRoot, f)))
-
-		const extension = vscode.extensions.getExtension<RooCodeAPI>("RooVeterinaryInc.roo-cline")
-
-		if (!extension) {
-			throw new Error("Extension not found")
-		}
-
-		const api = extension.isActive ? extension.exports : await extension.activate()
-
-		await api.setConfiguration({
-			apiProvider: "openrouter",
-			openRouterApiKey: process.env.OPENROUTER_API_KEY!,
-			openRouterModelId: "anthropic/claude-3.5-sonnet",
-		})
-
-		await waitUntilReady(api)
-
-		globalThis.api = api
-		globalThis.extension = extension
-
-		return new Promise<void>((resolve, reject) => {
-			try {
-				mocha.run((failures: number) => {
-					if (failures > 0) {
-						reject(new Error(`${failures} tests failed.`))
-					} else {
-						resolve()
-					}
-				})
-			} catch (err) {
-				console.error(err)
-				reject(err)
-			}
-		})
-	} catch (err) {
-		console.error("Error while running tests:")
-		console.error(err)
-		throw err
+	const extension = vscode.extensions.getExtension<RooCodeAPI>("RooVeterinaryInc.roo-cline")
+
+	if (!extension) {
+		throw new Error("Extension not found")
 	}
+
+	// Activate the extension if it's not already active.
+	const api = extension.isActive ? extension.exports : await extension.activate()
+
+	// TODO: We might want to support a "free" model out of the box so
+	// contributors can run the tests locally without having to pay.
+	await api.setConfiguration({
+		apiProvider: "openrouter",
+		openRouterApiKey: process.env.OPENROUTER_API_KEY!,
+		openRouterModelId: "anthropic/claude-3.5-sonnet",
+	})
+
+	await waitUntilReady({ api })
+
+	// Expose the API to the tests.
+	globalThis.api = api
+
+	// Add all the tests to the runner.
+	const mocha = new Mocha({ ui: "tdd", timeout: 300_000 })
+	const cwd = path.resolve(__dirname, "..")
+	;(await glob("**/**.test.js", { cwd })).forEach((testFile) => mocha.addFile(path.resolve(cwd, testFile)))
+
+	// Let's go!
+	return new Promise<void>((resolve, reject) =>
+		mocha.run((failures) => (failures === 0 ? resolve() : reject(new Error(`${failures} tests failed.`)))),
+	)
 }

e2e/src/suite/modes.test.ts

@@ -1,63 +1,34 @@
 import * as assert from "assert"
 
-import { waitForMessage } from "./utils"
+import { waitForMessage, getMessage } from "./utils"
 
 suite("Roo Code Modes", () => {
 	test("Should handle switching modes correctly", async function () {
-		const timeout = 300_000
 		const api = globalThis.api
 
-		const testPrompt =
+		let prompt =
 			"For each mode (Code, Architect, Ask) respond with the mode name and what it specializes in after switching to that mode, do not start with the current mode, be sure to say 'I AM DONE' after the task is complete."
 
 		await api.setConfiguration({ mode: "Code", alwaysAllowModeSwitch: true, autoApprovalEnabled: true })
-		await api.startNewTask(testPrompt)
+		let taskId = await api.startNewTask(prompt)
+		await waitForMessage({ api, taskId, include: "I AM DONE", exclude: "be sure to say", timeout: 300_000 })
 
-		await waitForMessage(api, { include: "I AM DONE", exclude: "be sure to say", timeout })
+		// Start grading portion of test to grade the response from 1 to 10.
+		prompt = `Given this prompt: ${prompt} grade the response from 1 to 10 in the format of "Grade: (1-10)": ${api
+			.getMessages(taskId)
+			.filter(({ type }) => type === "say")
+			.map(({ text }) => text ?? "")
+			.join("\n")}\nBe sure to say 'I AM DONE GRADING' after the task is complete.`
 
-		if (api.getMessages().length === 0) {
-			assert.fail("No messages received")
-		}
-
-		// Log the messages to the console.
-		api.getMessages().forEach(({ type, text }) => {
-			if (type === "say") {
-				console.log(text)
-			}
-		})
-
-		// Start Grading Portion of test to grade the response from 1 to 10.
 		await api.setConfiguration({ mode: "Ask" })
+		taskId = await api.startNewTask(prompt)
+		await waitForMessage({ api, taskId, include: "I AM DONE GRADING", exclude: "be sure to say" })
 
-		let output = api
-			.getMessages()
-			.map(({ type, text }) => (type === "say" ? text : ""))
-			.join("\n")
-
-		await api.startNewTask(
-			`Given this prompt: ${testPrompt} grade the response from 1 to 10 in the format of "Grade: (1-10)": ${output}\nBe sure to say 'I AM DONE GRADING' after the task is complete.`,
+		const match = getMessage({ api, taskId, include: "Grade:", exclude: "Grade: (1-10)" })?.text?.match(
+			/Grade: (\d+)/,
 		)
 
-		await waitForMessage(api, { include: "I AM DONE GRADING", exclude: "be sure to say", timeout })
-
-		if (api.getMessages().length === 0) {
-			assert.fail("No messages received")
-		}
-
-		api.getMessages().forEach(({ type, text }) => {
-			if (type === "say" && text?.includes("Grade:")) {
-				console.log(text)
-			}
-		})
-
-		const gradeMessage = api
-			.getMessages()
-			.find(
-				({ type, text }) => type === "say" && !text?.includes("Grade: (1-10)") && text?.includes("Grade:"),
-			)?.text
-
-		const gradeMatch = gradeMessage?.match(/Grade: (\d+)/)
-		const gradeNum = gradeMatch ? parseInt(gradeMatch[1]) : undefined
-		assert.ok(gradeNum !== undefined && gradeNum >= 7 && gradeNum <= 10, "Grade must be between 7 and 10")
+		const score = parseInt(match?.[1] ?? "0")
+		assert.ok(score >= 7 && score <= 10, "Grade must be between 7 and 10.")
 	})
 })

e2e/src/suite/subtasks.test.ts

@@ -1,58 +1,70 @@
 import * as assert from "assert"
 
-import { sleep, waitForToolUse, waitForMessage } from "./utils"
+import { sleep, waitForMessage, waitFor, getMessage } from "./utils"
 
 suite("Roo Code Subtasks", () => {
-	test.skip("Should handle subtask cancellation and resumption correctly", async function () {
+	test("Should handle subtask cancellation and resumption correctly", async function () {
 		const api = globalThis.api
 
 		await api.setConfiguration({
 			mode: "Code",
 			alwaysAllowModeSwitch: true,
 			alwaysAllowSubtasks: true,
 			autoApprovalEnabled: true,
+			enableCheckpoints: false,
 		})
 
+		const childPrompt = "You are a calculator. Respond only with numbers. What is the square root of 9?"
+
 		// Start a parent task that will create a subtask.
-		await api.startNewTask(
+		const parentTaskId = await api.startNewTask(
 			"You are the parent task. " +
-				"Create a subtask by using the new_task tool with the message 'You are the subtask'. " +
-				"After creating the subtask, wait for it to complete and then respond with 'Parent task resumed'.",
+				`Create a subtask by using the new_task tool with the message '${childPrompt}'.` +
+				"After creating the subtask, wait for it to complete and then respond 'Parent task resumed'.",
 		)
 
-		await waitForToolUse(api, "new_task")
+		let subTaskId: string | undefined = undefined
 
-		// Cancel the current task (which should be the subtask).
-		await api.cancelTask()
+		// Wait for the subtask to be spawned and then cancel it.
+		api.on("taskSpawned", (taskId) => (subTaskId = taskId))
+		await waitFor(() => !!subTaskId)
+		await sleep(2_000) // Give the task a chance to start and populate the history.
+		await api.cancelCurrentTask()
 
-		// Check if the parent task is still waiting (not resumed). We need to
-		// wait a bit to ensure any task resumption would have happened.
+		// Wait a bit to ensure any task resumption would have happened.
 		await sleep(5_000)
 
 		// The parent task should not have resumed yet, so we shouldn't see
 		// "Parent task resumed".
 		assert.ok(
-			!api.getMessages().some(({ type, text }) => type === "say" && text?.includes("Parent task resumed")),
-			"Parent task should not have resumed after subtask cancellation.",
+			getMessage({
+				api,
+				taskId: parentTaskId,
+				include: "Parent task resumed",
+				exclude: "You are the parent task",
+			}) === undefined,
+			"Parent task should not have resumed after subtask cancellation",
 		)
 
 		// Start a new task with the same message as the subtask.
-		await api.startNewTask("You are the subtask")
-
-		// Wait for the subtask to complete.
-		await waitForMessage(api, { include: "Task complete" })
+		const anotherTaskId = await api.startNewTask(childPrompt)
+		await waitForMessage({ taskId: anotherTaskId, api, include: "3" })
 
-		// Verify that the parent task is still not resumed. We need to wait a
-		// bit to ensure any task resumption would have happened.
+		// Wait a bit to ensure any task resumption would have happened.
 		await sleep(5_000)
 
 		// The parent task should still not have resumed.
 		assert.ok(
-			!api.getMessages().some(({ type, text }) => type === "say" && text?.includes("Parent task resumed")),
-			"Parent task should not have resumed after subtask completion.",
+			getMessage({
+				api,
+				taskId: parentTaskId,
+				include: "Parent task resumed",
+				exclude: "You are the parent task",
+			}) === undefined,
+			"Parent task should not have resumed after subtask cancellation",
 		)
 
 		// Clean up - cancel all tasks.
-		await api.cancelTask()
+		await api.cancelCurrentTask()
 	})
 })

e2e/src/suite/task.test.ts

@@ -4,7 +4,7 @@ suite("Roo Code Task", () => {
 	test("Should handle prompt and response correctly", async function () {
 		const api = globalThis.api
 		await api.setConfiguration({ mode: "Ask", alwaysAllowModeSwitch: true, autoApprovalEnabled: true })
-		await api.startNewTask("Hello world, what is your name? Respond with 'My name is ...'")
-		await waitForMessage(api, { include: "My name is Roo" })
+		const taskId = await api.startNewTask("Hello world, what is your name? Respond with 'My name is ...'")
+		await waitForMessage({ api, taskId, include: "My name is Roo" })
 	})
 })

e2e/src/suite/utils.ts

@@ -41,36 +41,51 @@ export const waitFor = (
 	])
 }
 
-export const waitUntilReady = async (api: RooCodeAPI, { timeout = 10_000, interval = 250 }: WaitForOptions = {}) => {
+type WaitUntilReadyOptions = WaitForOptions & {
+	api: RooCodeAPI
+}
+
+export const waitUntilReady = async ({ api, ...options }: WaitUntilReadyOptions) => {
 	await vscode.commands.executeCommand("roo-cline.SidebarProvider.focus")
-	await waitFor(api.isReady, { timeout, interval })
+	await waitFor(() => api.isReady(), options)
+}
+
+type WaitForToolUseOptions = WaitUntilReadyOptions & {
+	taskId: string
+	toolName: string
 }
 
-export const waitForToolUse = async (api: RooCodeAPI, toolName: string, options: WaitForOptions = {}) =>
+export const waitForToolUse = async ({ api, taskId, toolName, ...options }: WaitForToolUseOptions) =>
 	waitFor(
 		() =>
 			api
-				.getMessages()
+				.getMessages(taskId)
 				.some(({ type, say, text }) => type === "say" && say === "tool" && text && text.includes(toolName)),
 		options,
 	)
 
-export const waitForMessage = async (
-	api: RooCodeAPI,
-	options: WaitForOptions & { include: string; exclude?: string },
-) =>
-	waitFor(
-		() =>
-			api
-				.getMessages()
-				.some(
-					({ type, text }) =>
-						type === "say" &&
-						text &&
-						text.includes(options.include) &&
-						(!options.exclude || !text.includes(options.exclude)),
-				),
-		options,
-	)
+type WaitForMessageOptions = WaitUntilReadyOptions & {
+	taskId: string
+	include: string
+	exclude?: string
+}
+
+export const waitForMessage = async ({ api, taskId, include, exclude, ...options }: WaitForMessageOptions) =>
+	waitFor(() => !!getMessage({ api, taskId, include, exclude }), options)
+
+type GetMessageOptions = {
+	api: RooCodeAPI
+	taskId: string
+	include: string
+	exclude?: string
+}
+
+export const getMessage = ({ api, taskId, include, exclude }: GetMessageOptions) =>
+	api
+		.getMessages(taskId)
+		.find(
+			({ type, text }) =>
+				type === "say" && text && text.includes(include) && (!exclude || !text.includes(exclude)),
+		)
 
 export const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms))