From 9b56ad8c354714444d092782a29b7beede167238 Mon Sep 17 00:00:00 2001 From: cte Date: Mon, 31 Mar 2025 01:46:10 -0700 Subject: [PATCH] More progress --- benchmark/apps/cli/src/index.ts | 76 ++++++++++--------- benchmark/apps/web/src/app/runs/[id]/run.tsx | 4 +- .../web/src/app/runs/[id]/task-status.tsx | 15 ++-- .../apps/web/src/hooks/use-run-status.ts | 13 ++-- benchmark/packages/ipc/src/client.ts | 6 +- benchmark/packages/ipc/src/server.ts | 8 +- benchmark/packages/types/src/ipc.ts | 19 ++++- src/exports/api.ts | 45 ++++++++--- src/schemas/ipc.ts | 19 ++++- 9 files changed, 136 insertions(+), 69 deletions(-) diff --git a/benchmark/apps/cli/src/index.ts b/benchmark/apps/cli/src/index.ts index 29e6307d0d..f4fc8e38ea 100644 --- a/benchmark/apps/cli/src/index.ts +++ b/benchmark/apps/cli/src/index.ts @@ -95,23 +95,27 @@ const run = async (toolbox: GluegunToolbox) => { const server = new IpcServer(run.socketPath, () => {}) server.listen() - server.on(IpcMessageType.Connect, (clientId) => { - server.send(clientId, { - type: IpcMessageType.TaskEvent, - origin: IpcOrigin.Server, - // TODO: Broacast the set of running tasks. - data: { eventName: RooCodeEventName.Connect, taskId: -1 }, - }) - }) + // server.on(IpcMessageType.Connect, (clientId) => { + // server.send(clientId, { + // type: IpcMessageType.TaskEvent, + // origin: IpcOrigin.Server, + // data: { eventName: RooCodeEventName.Connect, taskId: -1 }, + // }) + // }) const chunks = inChunksOf(tasks, 3) for (const chunk of chunks) { await Promise.all( chunk.map(async (task) => { - const runSucceeded = await runExercise({ run, task, server }) - const passed = runSucceeded ? await runUnitTest({ task }) : false - await updateTask(task.id, { passed }) + if (task.finishedAt === null) { + await runExercise({ run, task, server }) + } + + if (task.passed === null) { + const passed = await runUnitTest({ task }) + await updateTask(task.id, { passed }) + } }), ) @@ -124,14 +128,7 @@ const run = async (toolbox: GluegunToolbox) => { const runExercise = async ({ run, task, server }: { run: Run; task: Task; server: IpcServer }) => { const { language, exercise } = task - - if (task.finishedAt) { - console.log(`Test result exists for ${language} / ${exercise}, skipping`) - return false - } - const prompt = fs.readFileSync(path.resolve(exercisesPath, `prompts/${language}.md`), "utf-8") - const dirname = path.dirname(run.socketPath) const taskSocketPath = path.resolve(dirname, `${dirname}/task-${task.id}.sock`) @@ -157,11 +154,10 @@ const runExercise = async ({ run, task, server }: { run: Run; task: Task; server if (!client.isReady) { client.disconnect() console.log(`[cli#runExercise | ${language} / ${exercise}] unable to connect`) - return false + return } let isTaskFinished = false - let isTaskAborted = false client.on(IpcMessageType.Disconnect, () => { console.log(`[cli#runExercise | ${language} / ${exercise}] disconnect`) @@ -176,6 +172,7 @@ const runExercise = async ({ run, task, server }: { run: Run; task: Task; server let taskStartedAt = Date.now() let taskMetricsId: number | undefined + let rooTaskId: string | undefined client.on(IpcMessageType.TaskEvent, async (taskEvent) => { const { eventName, payload } = taskEvent @@ -191,14 +188,6 @@ const runExercise = async ({ run, task, server }: { run: Run; task: Task; server console.log(`[cli#runExercise | ${language} / ${exercise}] taskEvent -> ${eventName}`) } - // if (eventName === RooCodeEventName.Message) { - // const { message: { ts, text, partial } } = payload[0] - - // if (!partial) { - // console.log(`${ts}: ${text}`) - // } - // } - if (eventName === RooCodeEventName.TaskStarted) { taskStartedAt = Date.now() @@ -216,6 +205,7 @@ const runExercise = async ({ run, task, server }: { run: Run; task: Task; server taskStartedAt = Date.now() taskMetricsId = taskMetrics.id + rooTaskId = payload[0] } if ( @@ -238,14 +228,10 @@ const runExercise = async ({ run, task, server }: { run: Run; task: Task; server }) } - if (eventName === RooCodeEventName.TaskCompleted) { + if (eventName === RooCodeEventName.TaskCompleted || eventName === RooCodeEventName.TaskAborted) { await updateTask(task.id, { finishedAt: new Date() }) isTaskFinished = true } - - if (eventName === RooCodeEventName.TaskAborted) { - isTaskAborted = true - } }) client.sendMessage({ @@ -269,13 +255,33 @@ const runExercise = async ({ run, task, server }: { run: Run; task: Task; server console.log(`[cli#runExercise | ${language} / ${exercise}] starting task`) try { - await pWaitFor(() => isTaskFinished || isTaskAborted, { interval: 1_000, timeout: 1 * 60 * 1_000 }) + await pWaitFor(() => isTaskFinished, { interval: 1_000, timeout: 1 * 60 * 1_000 }) // eslint-disable-next-line @typescript-eslint/no-unused-vars } catch (error) { console.log(`[cli#runExercise | ${language} / ${exercise}] time limit reached`) + + if (rooTaskId) { + client.sendMessage({ + type: IpcMessageType.TaskCommand, + origin: IpcOrigin.Client, + clientId: client.clientId!, + data: { commandName: TaskCommandName.CancelTask, data: rooTaskId }, + }) + + await new Promise((resolve) => setTimeout(resolve, 2_000)) + } + + await updateTask(task.id, { finishedAt: new Date() }) } try { + client.sendMessage({ + type: IpcMessageType.VSCodeCommand, + origin: IpcOrigin.Client, + clientId: client.clientId!, + data: "workbench.action.files.saveFiles", + }) + client.sendMessage({ type: IpcMessageType.VSCodeCommand, origin: IpcOrigin.Client, @@ -287,8 +293,6 @@ const runExercise = async ({ run, task, server }: { run: Run; task: Task; server } catch (error) { console.error(error) } - - return isTaskFinished } const runUnitTest = async ({ task }: { task: Task }) => { diff --git a/benchmark/apps/web/src/app/runs/[id]/run.tsx b/benchmark/apps/web/src/app/runs/[id]/run.tsx index 74ec301e4c..9a3fe3b520 100644 --- a/benchmark/apps/web/src/app/runs/[id]/run.tsx +++ b/benchmark/apps/web/src/app/runs/[id]/run.tsx @@ -25,7 +25,7 @@ import { TaskStatus } from "./task-status" import { ConnectionStatus } from "./connection-status" export function Run({ run }: { run: db.Run }) { - const { tasks, status, runningTaskId, output, outputCounts } = useRunStatus(run) + const { tasks, status, output, outputCounts } = useRunStatus(run) const scrollAreaRef = useRef(null) const [selectedTask, setSelectedTask] = useState() @@ -67,7 +67,7 @@ export function Run({ run }: { run: db.Run }) {
- +
{task.language}/{task.exercise}
diff --git a/benchmark/apps/web/src/app/runs/[id]/task-status.tsx b/benchmark/apps/web/src/app/runs/[id]/task-status.tsx index 40dffd7fec..0db97e2487 100644 --- a/benchmark/apps/web/src/app/runs/[id]/task-status.tsx +++ b/benchmark/apps/web/src/app/runs/[id]/task-status.tsx @@ -4,19 +4,18 @@ import { type Task } from "@benchmark/db" type TaskStatusProps = { task: Task - runningTaskId?: number } -export const TaskStatus = ({ task, runningTaskId }: TaskStatusProps) => { - return runningTaskId === task.id ? ( - - ) : !task.finishedAt ? ( - - ) : task.passed === false ? ( +export const TaskStatus = ({ task }: TaskStatusProps) => { + return task.passed === false ? ( ) : task.passed === true ? ( - ) : ( + ) : task.startedAt ? ( + ) : task.finishedAt ? ( + + ) : ( + ) } diff --git a/benchmark/apps/web/src/hooks/use-run-status.ts b/benchmark/apps/web/src/hooks/use-run-status.ts index 4e80b611d0..61d413bb94 100644 --- a/benchmark/apps/web/src/hooks/use-run-status.ts +++ b/benchmark/apps/web/src/hooks/use-run-status.ts @@ -8,12 +8,12 @@ import { getTasks } from "@/lib/server/tasks" import { useEventSource } from "@/hooks/use-event-source" export const useRunStatus = (run: Run) => { - const [runningTaskId, setRunningTaskId] = useState() + const [tasksUpdatedAt, setTasksUpdatedAt] = useState() const outputRef = useRef>(new Map()) const [outputCounts, setOutputCounts] = useState>({}) const { data: tasks } = useQuery({ - queryKey: ["run", run.id, runningTaskId], + queryKey: ["run", run.id, tasksUpdatedAt], queryFn: async () => getTasks(run.id), placeholderData: keepPreviousData, refetchInterval: 10_000, @@ -46,13 +46,10 @@ export const useRunStatus = (run: Run) => { } switch (eventName) { - case RooCodeEventName.Connect: - case RooCodeEventName.TaskCreated: case RooCodeEventName.TaskStarted: - setRunningTaskId(taskId) - break case RooCodeEventName.TaskCompleted: - setRunningTaskId(undefined) + case RooCodeEventName.TaskAborted: + setTasksUpdatedAt(Date.now()) break case RooCodeEventName.Message: { const [ @@ -79,5 +76,5 @@ export const useRunStatus = (run: Run) => { const status = useEventSource({ url, onMessage }) - return { tasks, status, runningTaskId, output: outputRef.current, outputCounts } + return { tasks, status, output: outputRef.current, outputCounts } } diff --git a/benchmark/packages/ipc/src/client.ts b/benchmark/packages/ipc/src/client.ts index 5b5bee988e..89728a5148 100644 --- a/benchmark/packages/ipc/src/client.ts +++ b/benchmark/packages/ipc/src/client.ts @@ -3,12 +3,12 @@ import * as crypto from "node:crypto" import ipc from "node-ipc" -import { IpcOrigin, IpcMessageType, IpcMessage, ipcMessageSchema, TaskCommand, TaskEvent } from "@benchmark/types" +import { IpcOrigin, IpcMessageType, IpcMessage, ipcMessageSchema, Ack, TaskCommand, TaskEvent } from "@benchmark/types" export type IpcClientEvents = { [IpcMessageType.Connect]: [] [IpcMessageType.Disconnect]: [] - [IpcMessageType.Ack]: [clientId: string] + [IpcMessageType.Ack]: [data: Ack] [IpcMessageType.TaskCommand]: [data: TaskCommand] [IpcMessageType.TaskEvent]: [data: TaskEvent] } @@ -75,7 +75,7 @@ export class IpcClient extends EventEmitter { switch (payload.type) { case IpcMessageType.Ack: this._clientId = payload.data.clientId - this.emit(IpcMessageType.Ack, payload.data.clientId) + this.emit(IpcMessageType.Ack, payload.data) break case IpcMessageType.TaskEvent: this.emit(IpcMessageType.TaskEvent, payload.data) diff --git a/benchmark/packages/ipc/src/server.ts b/benchmark/packages/ipc/src/server.ts index cb044f6fae..5faaae1766 100644 --- a/benchmark/packages/ipc/src/server.ts +++ b/benchmark/packages/ipc/src/server.ts @@ -47,7 +47,13 @@ export class IpcServer extends EventEmitter { const clientId = crypto.randomBytes(6).toString("hex") this._clients.set(clientId, socket) this.log(`[server#onConnect] clientId = ${clientId}, # clients = ${this._clients.size}`) - this.send(socket, { type: IpcMessageType.Ack, origin: IpcOrigin.Server, data: { clientId } }) + + this.send(socket, { + type: IpcMessageType.Ack, + origin: IpcOrigin.Server, + data: { clientId, pid: process.pid, ppid: process.ppid }, + }) + this.emit(IpcMessageType.Connect, clientId) } diff --git a/benchmark/packages/types/src/ipc.ts b/benchmark/packages/types/src/ipc.ts index cfc3c0c651..0efa5e1956 100644 --- a/benchmark/packages/types/src/ipc.ts +++ b/benchmark/packages/types/src/ipc.ts @@ -2,12 +2,25 @@ import { z } from "zod" import { RooCodeEventName, rooCodeEventsSchema, rooCodeSettingsSchema } from "./roo-code.js" +/** + * Ack + */ + +export const ackSchema = z.object({ + clientId: z.string(), + pid: z.number(), + ppid: z.number(), +}) + +export type Ack = z.infer + /** * TaskCommand */ export enum TaskCommandName { StartNewTask = "StartNewTask", + CancelTask = "CancelTask", } export const taskCommandSchema = z.discriminatedUnion("commandName", [ @@ -20,6 +33,10 @@ export const taskCommandSchema = z.discriminatedUnion("commandName", [ newTab: z.boolean().optional(), }), }), + z.object({ + commandName: z.literal(TaskCommandName.CancelTask), + data: z.string(), + }), ]) export type TaskCommand = z.infer @@ -110,7 +127,7 @@ export const ipcMessageSchema = z.discriminatedUnion("type", [ z.object({ type: z.literal(IpcMessageType.Ack), origin: z.literal(IpcOrigin.Server), - data: z.object({ clientId: z.string() }), + data: ackSchema, }), z.object({ type: z.literal(IpcMessageType.TaskCommand), diff --git a/src/exports/api.ts b/src/exports/api.ts index 8bbb93486e..be24483df3 100644 --- a/src/exports/api.ts +++ b/src/exports/api.ts @@ -16,6 +16,7 @@ export class API extends EventEmitter implements RooCodeAPI { private tabProvider?: ClineProvider private readonly context: vscode.ExtensionContext private readonly ipc?: IpcServer + private readonly taskMap = new Map() constructor(outputChannel: vscode.OutputChannel, provider: ClineProvider, socketPath?: string) { super() @@ -39,7 +40,10 @@ export class API extends EventEmitter implements RooCodeAPI { switch (commandName) { case TaskCommandName.StartNewTask: - this.startNewTask(data) + await this.startNewTask(data) + break + case TaskCommandName.CancelTask: + await this.cancelTask(data) break } }) @@ -102,7 +106,6 @@ export class API extends EventEmitter implements RooCodeAPI { await provider.postMessageToWebview({ type: "invoke", invoke: "newChat", text, images }) const { taskId } = await provider.initClineWithTask(text, images) - return taskId } @@ -114,6 +117,15 @@ export class API extends EventEmitter implements RooCodeAPI { await this.sidebarProvider.finishSubTask(lastMessage) } + public async cancelTask(taskId: string) { + const provider = this.taskMap.get(taskId) + + if (provider) { + await provider.cancelTask() + this.taskMap.delete(taskId) + } + } + public async cancelCurrentTask() { await this.sidebarProvider.cancelTask() } @@ -156,18 +168,33 @@ export class API extends EventEmitter implements RooCodeAPI { private registerListeners(provider: ClineProvider) { provider.on("clineCreated", (cline) => { + cline.on("taskStarted", () => { + this.emit(RooCodeEventName.TaskStarted, cline.taskId) + this.taskMap.set(cline.taskId, provider) + }) + cline.on("message", (message) => this.emit(RooCodeEventName.Message, { taskId: cline.taskId, ...message })) - cline.on("taskStarted", () => this.emit(RooCodeEventName.TaskStarted, cline.taskId)) - cline.on("taskPaused", () => this.emit(RooCodeEventName.TaskPaused, cline.taskId)) - cline.on("taskUnpaused", () => this.emit(RooCodeEventName.TaskUnpaused, cline.taskId)) - cline.on("taskAskResponded", () => this.emit(RooCodeEventName.TaskAskResponded, cline.taskId)) - cline.on("taskAborted", () => this.emit(RooCodeEventName.TaskAborted, cline.taskId)) - cline.on("taskSpawned", (childTaskId) => this.emit(RooCodeEventName.TaskSpawned, cline.taskId, childTaskId)) - cline.on("taskCompleted", (_, usage) => this.emit(RooCodeEventName.TaskCompleted, cline.taskId, usage)) + cline.on("taskTokenUsageUpdated", (_, usage) => this.emit(RooCodeEventName.TaskTokenUsageUpdated, cline.taskId, usage), ) + cline.on("taskAskResponded", () => this.emit(RooCodeEventName.TaskAskResponded, cline.taskId)) + + cline.on("taskAborted", () => { + this.emit(RooCodeEventName.TaskAborted, cline.taskId) + this.taskMap.delete(cline.taskId) + }) + + cline.on("taskCompleted", (_, usage) => { + this.emit(RooCodeEventName.TaskCompleted, cline.taskId, usage) + this.taskMap.delete(cline.taskId) + }) + + cline.on("taskSpawned", (childTaskId) => this.emit(RooCodeEventName.TaskSpawned, cline.taskId, childTaskId)) + cline.on("taskPaused", () => this.emit(RooCodeEventName.TaskPaused, cline.taskId)) + cline.on("taskUnpaused", () => this.emit(RooCodeEventName.TaskUnpaused, cline.taskId)) + this.emit(RooCodeEventName.TaskCreated, cline.taskId) }) } diff --git a/src/schemas/ipc.ts b/src/schemas/ipc.ts index 70eb93648a..06b2ac64dc 100644 --- a/src/schemas/ipc.ts +++ b/src/schemas/ipc.ts @@ -2,12 +2,25 @@ import { z } from "zod" import { RooCodeEventName, rooCodeEventsSchema, rooCodeSettingsSchema } from "./index" +/** + * Ack + */ + +export const ackSchema = z.object({ + clientId: z.string(), + pid: z.number(), + ppid: z.number(), +}) + +export type Ack = z.infer + /** * TaskCommand */ export enum TaskCommandName { StartNewTask = "StartNewTask", + CancelTask = "CancelTask", } export const taskCommandSchema = z.discriminatedUnion("commandName", [ @@ -20,6 +33,10 @@ export const taskCommandSchema = z.discriminatedUnion("commandName", [ newTab: z.boolean().optional(), }), }), + z.object({ + commandName: z.literal(TaskCommandName.CancelTask), + data: z.string(), + }), ]) export type TaskCommand = z.infer @@ -95,7 +112,7 @@ export const ipcMessageSchema = z.discriminatedUnion("type", [ z.object({ type: z.literal(IpcMessageType.Ack), origin: z.literal(IpcOrigin.Server), - data: z.object({ clientId: z.string(), pid: z.number(), ppid: z.number() }), + data: ackSchema, }), z.object({ type: z.literal(IpcMessageType.TaskCommand),