From 8f64109894f8a884f13154dbde18955673d51488 Mon Sep 17 00:00:00 2001
From: cte <cestreich@gmail.com>
Date: Mon, 24 Mar 2025 22:55:00 -0700
Subject: [PATCH] Improve task progress UI

---
 .../src/app/runs/[id]/connection-status.tsx   |  4 +-
 benchmark/apps/web/src/app/runs/[id]/run.tsx  | 96 ++++++++++++++-----
 benchmark/apps/web/src/lib/format-tokens.ts   |  7 ++
 benchmark/apps/web/src/lib/index.ts           |  1 +
 benchmark/packages/db/src/queries/tasks.ts    |  3 +-
 5 files changed, 85 insertions(+), 26 deletions(-)
 create mode 100644 benchmark/apps/web/src/lib/format-tokens.ts

diff --git a/benchmark/apps/web/src/app/runs/[id]/connection-status.tsx b/benchmark/apps/web/src/app/runs/[id]/connection-status.tsx
index 95c069f7fd..f1c39ce03c 100644
--- a/benchmark/apps/web/src/app/runs/[id]/connection-status.tsx
+++ b/benchmark/apps/web/src/app/runs/[id]/connection-status.tsx
@@ -3,18 +3,16 @@ import { cn } from "@/lib/utils"
 
 type ConnectionStatusProps = {
 	status: EventSourceStatus
-	clientId?: string
 	pid: number | null
 }
 
-export const ConnectionStatus = ({ status, clientId, pid }: ConnectionStatusProps) => (
+export const ConnectionStatus = ({ status, pid }: ConnectionStatusProps) => (
 	<div className="flex items-center">
 		<div className="flex flex-col items-end gap-1 font-mono text-xs border-r border-dotted pr-4 mr-4">
 			<div>
 				Status: <span className="capitalize">{status}</span>
 			</div>
 			<div>PID: {pid}</div>
-			{clientId && <div>Client: {clientId}</div>}
 		</div>
 		<div className="relative">
 			<div
diff --git a/benchmark/apps/web/src/app/runs/[id]/run.tsx b/benchmark/apps/web/src/app/runs/[id]/run.tsx
index cc689e8a33..0cb1590df1 100644
--- a/benchmark/apps/web/src/app/runs/[id]/run.tsx
+++ b/benchmark/apps/web/src/app/runs/[id]/run.tsx
@@ -1,18 +1,31 @@
 "use client"
 
 import { useState, useRef, useEffect } from "react"
-import { LoaderCircle, RectangleEllipsis } from "lucide-react"
+import { LoaderCircle, SquareTerminal } from "lucide-react"
 
 import * as db from "@benchmark/db"
 
+import { formatCurrency, formatDuration, formatTokens } from "@/lib"
 import { useRunStatus } from "@/hooks/use-run-status"
-import { Drawer, DrawerContent, DrawerHeader, DrawerTitle, ScrollArea } from "@/components/ui"
+import {
+	Drawer,
+	DrawerContent,
+	DrawerHeader,
+	DrawerTitle,
+	ScrollArea,
+	Table,
+	TableBody,
+	TableCell,
+	TableHead,
+	TableHeader,
+	TableRow,
+} from "@/components/ui"
 
 import { TaskStatus } from "./task-status"
 import { ConnectionStatus } from "./connection-status"
 
 export function Run({ run }: { run: db.Run }) {
-	const { tasks, status, clientId, runningTaskId, output, outputCounts } = useRunStatus(run)
+	const { tasks, status, runningTaskId, output, outputCounts } = useRunStatus(run)
 	const scrollAreaRef = useRef<HTMLDivElement>(null)
 	const [selectedTask, setSelectedTask] = useState<db.Task>()
 
@@ -32,34 +45,73 @@ export function Run({ run }: { run: db.Run }) {
 	return (
 		<>
 			<div className="flex flex-col gap-2">
-				<div className="border-b mb-2 pb-2">
-					<div>Run #{run.id}</div>
+				<div>
 					<div>{run.model}</div>
 					{run.description && <div className="text-sm text-muted-foreground">{run.description}</div>}
 				</div>
 				{!tasks ? (
 					<LoaderCircle className="size-4 animate-spin" />
 				) : (
-					tasks.map((task) => (
-						<div key={task.id} className="flex items-center gap-2">
-							<TaskStatus task={task} runningTaskId={runningTaskId} />
-							<div>
-								{task.language}/{task.exercise}
-							</div>
-							{(outputCounts[task.id] ?? 0) > 0 && (
-								<div
-									className="flex items-center gap-1 cursor-pointer"
-									onClick={() => setSelectedTask(task)}>
-									<RectangleEllipsis className="size-4" />
-									<div className="text-xs">({outputCounts[task.id]})</div>
-								</div>
-							)}
-						</div>
-					))
+					<Table className="border">
+						<TableHeader>
+							<TableRow>
+								<TableHead>Exercise</TableHead>
+								<TableHead className="text-center">Tokens In / Out</TableHead>
+								<TableHead>Context</TableHead>
+								<TableHead>Duration</TableHead>
+								<TableHead>Cost</TableHead>
+							</TableRow>
+						</TableHeader>
+						<TableBody>
+							{tasks.map((task) => (
+								<TableRow key={task.id}>
+									<TableCell>
+										<div className="flex items-center gap-2">
+											<TaskStatus task={task} runningTaskId={runningTaskId} />
+											<div>
+												{task.language}/{task.exercise}
+											</div>
+											{(outputCounts[task.id] ?? 0) > 0 && (
+												<div
+													className="flex items-center gap-1 cursor-pointer"
+													onClick={() => setSelectedTask(task)}>
+													<SquareTerminal className="size-4" />
+													<div className="font-mono text-xs text-foreground/50">
+														{outputCounts[task.id]}
+													</div>
+												</div>
+											)}
+										</div>
+									</TableCell>
+									{task.taskMetrics ? (
+										<>
+											<TableCell className="font-mono text-xs">
+												<div className="flex items-center justify-evenly">
+													<div>{formatTokens(task.taskMetrics.tokensIn)}</div>/
+													<div>{formatTokens(task.taskMetrics.tokensOut)}</div>
+												</div>
+											</TableCell>
+											<TableCell className="font-mono text-xs">
+												{formatTokens(task.taskMetrics.tokensContext)}
+											</TableCell>
+											<TableCell className="font-mono text-xs">
+												{formatDuration(task.taskMetrics.duration)}
+											</TableCell>
+											<TableCell className="font-mono text-xs">
+												{formatCurrency(task.taskMetrics.cost)}
+											</TableCell>
+										</>
+									) : (
+										<TableCell colSpan={3} />
+									)}
+								</TableRow>
+							))}
+						</TableBody>
+					</Table>
 				)}
 			</div>
 			<div className="absolute top-5 right-5">
-				<ConnectionStatus status={status} clientId={clientId} pid={run.pid} />
+				<ConnectionStatus status={status} pid={run.pid} />
 			</div>
 			<Drawer open={!!selectedTask} onOpenChange={() => setSelectedTask(undefined)}>
 				<DrawerContent>
diff --git a/benchmark/apps/web/src/lib/format-tokens.ts b/benchmark/apps/web/src/lib/format-tokens.ts
new file mode 100644
index 0000000000..d017c9ce6a
--- /dev/null
+++ b/benchmark/apps/web/src/lib/format-tokens.ts
@@ -0,0 +1,7 @@
+export const formatTokens = (tokens: number) => {
+	if (tokens < 1000) {
+		return tokens.toString()
+	}
+
+	return `${(tokens / 1000).toFixed(1)}k`
+}
diff --git a/benchmark/apps/web/src/lib/index.ts b/benchmark/apps/web/src/lib/index.ts
index e2d57baac5..f4262c384f 100644
--- a/benchmark/apps/web/src/lib/index.ts
+++ b/benchmark/apps/web/src/lib/index.ts
@@ -1,2 +1,3 @@
 export { formatCurrency } from "./format-currency"
 export { formatDuration } from "./format-duration"
+export { formatTokens } from "./format-tokens"
diff --git a/benchmark/packages/db/src/queries/tasks.ts b/benchmark/packages/db/src/queries/tasks.ts
index 77ea4f678d..a32e92bf51 100644
--- a/benchmark/packages/db/src/queries/tasks.ts
+++ b/benchmark/packages/db/src/queries/tasks.ts
@@ -58,4 +58,5 @@ export const getTask = async ({ runId, language, exercise }: GetTask) =>
 		where: and(eq(table.runId, runId), eq(table.language, language), eq(table.exercise, exercise)),
 	})
 
-export const getTasks = async (runId: number) => db.query.tasks.findMany({ where: eq(table.runId, runId) })
+export const getTasks = async (runId: number) =>
+	db.query.tasks.findMany({ where: eq(table.runId, runId), with: { taskMetrics: true } })