diff --git a/benchmark/apps/web/src/app/runs/[id]/schemas.ts b/benchmark/apps/web/src/app/runs/[id]/schemas.ts index 91aaa579ff..8fc15181a7 100644 --- a/benchmark/apps/web/src/app/runs/[id]/schemas.ts +++ b/benchmark/apps/web/src/app/runs/[id]/schemas.ts @@ -1,42 +1,20 @@ import { z } from "zod" -export const messageSchema = z.object({ - type: z.enum(["hello", "data"]), - data: z.record(z.string(), z.unknown()), -}) - -export const taskSchema = z.object({ - id: z.number(), -}) - -export const taskEventSchema = z.discriminatedUnion("event", [ - z.object({ event: z.literal("client"), task: taskSchema }), - z.object({ event: z.literal("taskStarted"), task: taskSchema }), - z.object({ event: z.literal("taskFinished"), task: taskSchema }), +export const ipcServerMessageSchema = z.discriminatedUnion("type", [ z.object({ - event: z.literal("message"), - task: taskSchema, - message: z.object({ - taskId: z.string(), - action: z.enum(["created", "updated"]), - message: z.object({ - ts: z.number(), - type: z.enum(["ask", "say"]), - text: z.string(), - partial: z.boolean().optional(), + type: z.literal("Ack"), + data: z.object({ clientId: z.string() }), + }), + z.object({ + type: z.literal("TaskEvent"), + data: z.object({ + eventName: z.enum(["connect", "taskStarted", "message", "taskTokenUsageUpdated", "taskFinished"]), + data: z.object({ + task: z.object({ id: z.number() }), + // message: z.object({}).optional(), + // usage: z.object({}).optional(), + // taskMetrics: z.object({}).optional(), }), }), }), - z.object({ - event: z.literal("taskTokenUsageUpdated"), - task: taskSchema, - usage: z.object({ - totalTokensIn: z.number(), - totalTokensOut: z.number(), - totalCacheWrites: z.number().optional(), - totalCacheReads: z.number().optional(), - totalCost: z.number(), - contextTokens: z.number(), - }), - }), ]) diff --git a/benchmark/apps/web/src/app/runs/[id]/use-run-status.ts b/benchmark/apps/web/src/app/runs/[id]/use-run-status.ts index c9724e4751..344b6efca6 100644 --- a/benchmark/apps/web/src/app/runs/[id]/use-run-status.ts +++ b/benchmark/apps/web/src/app/runs/[id]/use-run-status.ts @@ -1,12 +1,12 @@ -import { useState } from "react" -import { keepPreviousData, useQuery } from "@tanstack/react-query" -import { useCallback } from "react" +import { useState, useCallback } from "react" +import { useQuery, keepPreviousData } from "@tanstack/react-query" + import { useEventSource } from "@/hooks/use-event-source" import { Run } from "@benchmark/db" import { getTasks } from "./actions" -import { messageSchema, taskEventSchema } from "./schemas" +import { ipcServerMessageSchema } from "./schemas" export const useRunStatus = (run: Run) => { const [clientId, setClientId] = useState() @@ -30,7 +30,7 @@ export const useRunStatus = (run: Run) => { return } - const result = messageSchema.safeParse(data) + const result = ipcServerMessageSchema.safeParse(data) if (!result.success) { console.log(`unrecognized messageEvent.data: ${messageEvent.data}`) @@ -39,26 +39,17 @@ export const useRunStatus = (run: Run) => { const payload = result.data - if (payload.type === "hello") { + if (payload.type === "Ack") { setClientId(payload.data.clientId as string) - } else if (payload.type === "data") { - const taskEvent = taskEventSchema.safeParse(payload.data) + } else if (payload.type === "TaskEvent") { + const { + eventName, + data: { task }, + } = payload.data - if (!taskEvent.success) { - console.log(`unrecognized payload.data`, payload.data, taskEvent.error) - return - } - - if (taskEvent.data.event === "client") { - console.log(`client`, taskEvent.data) - setRunningTaskId(taskEvent.data.task.id) - } else if (taskEvent.data.event === "message") { - // console.log(`message: ${taskEvent.data.message.message.text}`) - } else if (taskEvent.data.event === "taskTokenUsageUpdated") { - console.log(`taskTokenUsageUpdated`, taskEvent.data.usage) - } else if (taskEvent.data.event === "taskStarted") { - setRunningTaskId(taskEvent.data.task.id) - } else if (taskEvent.data.event === "taskFinished") { + if (eventName === "connect" || eventName === "taskStarted") { + setRunningTaskId(task.id) + } else if (eventName === "taskFinished") { setRunningTaskId(undefined) } }