Fix schemas

Author: cte

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({
-		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(),
+		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(),
+			}),
 		}),
 	}),
 ])

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<string>()
@@ -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)
-
-			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") {
+		} else if (payload.type === "TaskEvent") {
+			const {
+				eventName,
+				data: { task },
+			} = payload.data
+
+			if (eventName === "connect" || eventName === "taskStarted") {
+				setRunningTaskId(task.id)
+			} else if (eventName === "taskFinished") {
 				setRunningTaskId(undefined)
 			}
 		}