|
1 | | -import { useState, useEffect, useRef } from "react"; |
| 1 | +import { useEffect, useSyncExternalStore } from "react"; |
2 | 2 | import { useChatStore, useSettingsStore } from "@/stores"; |
3 | 3 | import { cn } from "@/lib/utils"; |
4 | 4 | import { IconArrowUp, IconArrowDown, IconGauge, IconRefresh, IconBolt } from "@tabler/icons-react"; |
5 | 5 | import { Tooltip, TooltipContent, TooltipTrigger } from "./ui/tooltip"; |
6 | 6 |
|
| 7 | +/** |
| 8 | + * Generation speed is one property of one stream, so it is measured once here |
| 9 | + * rather than per component. |
| 10 | + * |
| 11 | + * This used to be per-hook state. `TokenInfo` and `ChatStatus` are mounted at |
| 12 | + * the same time and both call it, so each kept its own start timestamp and |
| 13 | + * token baseline and averaged over a different window — the header and the |
| 14 | + * expanded panel showed different numbers for the same response. |
| 15 | + * |
| 16 | + * The 250ms sampler also cannot depend on the token count: tokens change on |
| 17 | + * every streamed chunk, and re-running the effect per chunk tore down the |
| 18 | + * interval before it could fire, pinning the readout at 0 for fast streams. |
| 19 | + */ |
| 20 | +const speedListeners = new Set<() => void>(); |
| 21 | +let currentSpeed = 0; |
| 22 | +let sampler: ReturnType<typeof setInterval> | null = null; |
| 23 | +let startedAt: number | null = null; |
| 24 | +let startTokens = 0; |
| 25 | +let latestTokens = 0; |
| 26 | + |
| 27 | +function setSpeed(next: number): void { |
| 28 | + if (next === currentSpeed) return; |
| 29 | + currentSpeed = next; |
| 30 | + for (const listener of speedListeners) listener(); |
| 31 | +} |
| 32 | + |
| 33 | +function beginMeasuring(tokens: number): void { |
| 34 | + startedAt = Date.now(); |
| 35 | + startTokens = tokens; |
| 36 | + latestTokens = tokens; |
| 37 | + sampler ??= setInterval(() => { |
| 38 | + if (startedAt === null) return; |
| 39 | + const elapsedSec = (Date.now() - startedAt) / 1000; |
| 40 | + const generated = latestTokens - startTokens; |
| 41 | + if (elapsedSec > 0.2 && generated >= 0) setSpeed(generated / elapsedSec); |
| 42 | + }, 250); |
| 43 | +} |
| 44 | + |
| 45 | +function stopMeasuring(): void { |
| 46 | + startedAt = null; |
| 47 | + if (sampler !== null) { |
| 48 | + clearInterval(sampler); |
| 49 | + sampler = null; |
| 50 | + } |
| 51 | + // A finished stream has no rate; leaving the last value up made a stale |
| 52 | + // number look live. |
| 53 | + setSpeed(0); |
| 54 | +} |
| 55 | + |
| 56 | +function subscribeToSpeed(listener: () => void): () => void { |
| 57 | + speedListeners.add(listener); |
| 58 | + return () => speedListeners.delete(listener); |
| 59 | +} |
| 60 | + |
7 | 61 | export function useTokenSpeed() { |
8 | 62 | const isStreaming = useChatStore((s) => s.isStreaming); |
9 | 63 | const outputTokens = useChatStore((s) => s.tokenUsage.output + s.activeTokenUsage.output); |
| 64 | + const speed = useSyncExternalStore(subscribeToSpeed, () => currentSpeed); |
10 | 65 |
|
11 | | - const [speed, setSpeed] = useState<number>(0); |
12 | | - const startTimeRef = useRef<number | null>(null); |
13 | | - const startTokensRef = useRef<number>(0); |
| 66 | + latestTokens = outputTokens; |
14 | 67 |
|
15 | 68 | useEffect(() => { |
16 | | - if (isStreaming) { |
17 | | - if (startTimeRef.current === null) { |
18 | | - startTimeRef.current = Date.now(); |
19 | | - startTokensRef.current = outputTokens; |
20 | | - } |
21 | | - |
22 | | - const interval = setInterval(() => { |
23 | | - if (!startTimeRef.current) return; |
24 | | - const elapsedSec = (Date.now() - startTimeRef.current) / 1000; |
25 | | - const tokensGenerated = outputTokens - startTokensRef.current; |
26 | | - if (elapsedSec > 0.2 && tokensGenerated >= 0) { |
27 | | - setSpeed(tokensGenerated / elapsedSec); |
28 | | - } |
29 | | - }, 250); |
30 | | - |
31 | | - return () => clearInterval(interval); |
| 69 | + if (!isStreaming) { |
| 70 | + stopMeasuring(); |
| 71 | + return; |
32 | 72 | } |
33 | | - startTimeRef.current = null; |
34 | | - |
35 | | - }, [isStreaming, outputTokens]); |
| 73 | + // Idempotent across concurrent consumers: the first one to see the stream |
| 74 | + // start defines the window, the rest attach to the same measurement. |
| 75 | + if (startedAt === null) beginMeasuring(latestTokens); |
| 76 | + }, [isStreaming]); |
36 | 77 |
|
37 | 78 | return { speed, isStreaming }; |
38 | 79 | } |
@@ -78,7 +119,7 @@ export function TokenInfo() { |
78 | 119 | <span className="text-muted-foreground text-[10px] flex items-center gap-1"> |
79 | 120 | <IconBolt className="size-3 text-amber-400" /> Generation Speed |
80 | 121 | </span> |
81 | | - <span className="text-sm font-semibold font-mono text-foreground mt-0.5"> |
| 122 | + <span className="text-sm font-semibold font-mono text-foreground mt-0.5 whitespace-nowrap tabular-nums"> |
82 | 123 | {speed > 0 ? `${speed.toFixed(1)} tok/s` : "Idle"} |
83 | 124 | </span> |
84 | 125 | <span className="text-[10px] text-muted-foreground font-mono mt-0.5"> |
@@ -150,9 +191,13 @@ export function ChatStatus() { |
150 | 191 | {speed > 0 && ( |
151 | 192 | <Tooltip> |
152 | 193 | <TooltipTrigger asChild> |
153 | | - <span className="flex items-center gap-1 text-amber-500 font-mono font-medium"> |
154 | | - <IconBolt className="size-3 fill-amber-400/20 text-amber-400 animate-pulse" /> |
155 | | - <span>{speed.toFixed(1)} t/s</span> |
| 194 | + <span className="flex items-center gap-1 whitespace-nowrap text-amber-500 font-mono font-medium"> |
| 195 | + <IconBolt className="size-3 shrink-0 fill-amber-400/20 text-amber-400 animate-pulse" /> |
| 196 | + {/* tabular-nums keeps the pill from resizing as the rate changes, |
| 197 | + and nowrap stops "74.7" and "t/s" breaking onto two lines in a |
| 198 | + narrow sidebar. */} |
| 199 | + <span className="tabular-nums">{speed.toFixed(1)}</span> |
| 200 | + <span>t/s</span> |
156 | 201 | </span> |
157 | 202 | </TooltipTrigger> |
158 | 203 | <TooltipContent>Live Generation Speed (tokens / second)</TooltipContent> |
|
0 commit comments