Fix schemas

This commit is contained in:
cte 2025-03-22 02:24:01 -07:00
parent 39b300e82d
commit 0a9fbe68c2
2 changed files with 27 additions and 58 deletions

View file

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

View file

@ -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)
} 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)
}
}