Skip to content

Commit 6bd2503

Browse files
ashtonandersonashtonanderson
authored
Merge pull request #233 from CSSLab/codex/stockfish-loading-cache-hardening
Codex/stockfish loading cache hardening
2 parents 72ca23e + a07d3e9 commit 6bd2503

5 files changed

Lines changed: 351 additions & 72 deletions

File tree

‎src/contexts/MaiaEngineContext.tsx‎

Lines changed: 56 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ import React, {
1010
} from 'react'
1111
import toast from 'react-hot-toast'
1212

13+
const MAIA_LOADING_TOAST_DELAY_MS = 800
14+
1315
export const MaiaEngineContext = React.createContext<MaiaEngine>({
1416
maia: undefined,
1517
status: 'loading',
@@ -28,6 +30,7 @@ export const MaiaEngineContextProvider: React.FC<{ children: ReactNode }> = ({
2830
const [progress, setProgress] = useState(0)
2931
const [error, setError] = useState<string | null>(null)
3032
const toastId = useRef<string | null>(null)
33+
const loadingToastTimerRef = useRef<number | null>(null)
3134

3235
const maia = useMemo(() => {
3336
const model = new Maia({
@@ -63,33 +66,77 @@ export const MaiaEngineContextProvider: React.FC<{ children: ReactNode }> = ({
6366
// Toast notifications for Maia model status
6467
useEffect(() => {
6568
return () => {
66-
toast.dismiss()
69+
if (
70+
loadingToastTimerRef.current !== null &&
71+
typeof window !== 'undefined'
72+
) {
73+
window.clearTimeout(loadingToastTimerRef.current)
74+
loadingToastTimerRef.current = null
75+
}
76+
if (toastId.current) {
77+
toast.dismiss(toastId.current)
78+
toastId.current = null
79+
}
6780
}
6881
}, [])
6982

7083
useEffect(() => {
71-
if (status === 'loading' && !toastId.current) {
72-
toastId.current = toast.loading('Loading Maia Model...')
73-
} else if (status === 'ready') {
84+
if (status === 'loading') {
85+
if (
86+
typeof window === 'undefined' ||
87+
toastId.current ||
88+
loadingToastTimerRef.current !== null
89+
) {
90+
return
91+
}
92+
93+
loadingToastTimerRef.current = window.setTimeout(() => {
94+
loadingToastTimerRef.current = null
95+
if (!toastId.current) {
96+
toastId.current = toast.loading('Loading Maia...')
97+
}
98+
}, MAIA_LOADING_TOAST_DELAY_MS)
99+
return
100+
}
101+
102+
if (
103+
loadingToastTimerRef.current !== null &&
104+
typeof window !== 'undefined'
105+
) {
106+
window.clearTimeout(loadingToastTimerRef.current)
107+
loadingToastTimerRef.current = null
108+
}
109+
110+
if (status === 'no-cache' || status === 'downloading') {
111+
if (toastId.current) {
112+
toast.dismiss(toastId.current)
113+
toastId.current = null
114+
}
115+
return
116+
}
117+
118+
if (status === 'ready') {
119+
// Only show success if a loading toast was visible.
74120
if (toastId.current) {
75121
toast.success('Loaded Maia! Analysis is ready', {
76122
id: toastId.current,
77123
})
78124
toastId.current = null
79-
} else {
80-
toast.success('Loaded Maia! Analysis is ready')
81125
}
82126
} else if (status === 'error') {
127+
const message = error
128+
? `Failed to load Maia model: ${error}`
129+
: 'Failed to load Maia model'
83130
if (toastId.current) {
84-
toast.error('Failed to load Maia model', {
131+
toast.error(message, {
85132
id: toastId.current,
86133
})
87134
toastId.current = null
88135
} else {
89-
toast.error('Failed to load Maia model')
136+
toast.error(message)
90137
}
91138
}
92-
}, [status])
139+
}, [status, error])
93140

94141
return (
95142
<MaiaEngineContext.Provider

‎src/contexts/StockfishEngineContext.tsx‎

Lines changed: 148 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,72 @@ import toast from 'react-hot-toast'
1010
import { StockfishStatus, StockfishEngine } from 'src/types'
1111
import Engine from 'src/lib/engine/stockfish'
1212

13+
const STOCKFISH_LOADING_TOAST_DELAY_MS = 800
14+
const STOCKFISH_DEBUG_LOADING_KEY = 'maia.stockfishDebugLoading'
15+
16+
const isTruthy = (value: string | null | undefined): boolean => {
17+
if (!value) return false
18+
return ['1', 'true', 'yes', 'on'].includes(value.toLowerCase())
19+
}
20+
21+
const isStockfishDebugLoadingEnabled = (): boolean => {
22+
if (typeof window !== 'undefined') {
23+
try {
24+
const localValue = window.localStorage.getItem(
25+
STOCKFISH_DEBUG_LOADING_KEY,
26+
)
27+
if (localValue !== null) return isTruthy(localValue)
28+
} catch {
29+
// ignore localStorage access failures
30+
}
31+
}
32+
33+
return isTruthy(process.env.NEXT_PUBLIC_STOCKFISH_DEBUG_LOADING)
34+
}
35+
36+
const getStockfishLoadingLabel = (
37+
engine: Engine | null,
38+
debugLoadingEnabled: boolean,
39+
): string => {
40+
if (!debugLoadingEnabled) {
41+
return 'Loading Stockfish...'
42+
}
43+
44+
if (!engine) return 'Loading Stockfish...'
45+
46+
switch (engine.initializationPhase) {
47+
case 'checking-cache':
48+
return 'Checking local Stockfish cache...'
49+
case 'downloading-nnue':
50+
return 'Downloading Stockfish model weights...'
51+
case 'loading-nnue':
52+
return 'Loading Stockfish from local cache...'
53+
case 'loading-module':
54+
return 'Starting Stockfish engine...'
55+
default:
56+
return 'Loading Stockfish...'
57+
}
58+
}
59+
60+
let sharedClientStockfishEngine: Engine | null = null
61+
62+
const getOrCreateStockfishEngine = (): Engine => {
63+
if (typeof window === 'undefined') {
64+
return new Engine()
65+
}
66+
67+
if (
68+
!sharedClientStockfishEngine ||
69+
(sharedClientStockfishEngine.initializationError &&
70+
!sharedClientStockfishEngine.ready &&
71+
!sharedClientStockfishEngine.initializing)
72+
) {
73+
sharedClientStockfishEngine = new Engine()
74+
}
75+
76+
return sharedClientStockfishEngine
77+
}
78+
1379
export const StockfishEngineContext = React.createContext<StockfishEngine>({
1480
streamEvaluations: () => {
1581
throw new Error(
@@ -32,13 +98,24 @@ export const StockfishEngineContextProvider: React.FC<{
3298
children: ReactNode
3399
}> = ({ children }: { children: ReactNode }) => {
34100
const engineRef = useRef<Engine | null>(null)
35-
const [status, setStatus] = useState<StockfishStatus>('loading')
36-
const [error, setError] = useState<string | null>(null)
37-
const toastId = useRef<string | null>(null)
38-
39101
if (!engineRef.current) {
40-
engineRef.current = new Engine()
102+
engineRef.current = getOrCreateStockfishEngine()
41103
}
104+
const [status, setStatus] = useState<StockfishStatus>(() => {
105+
if (engineRef.current?.initializationError) return 'error'
106+
return engineRef.current?.ready ? 'ready' : 'loading'
107+
})
108+
const [error, setError] = useState<string | null>(
109+
() => engineRef.current?.initializationError ?? null,
110+
)
111+
const [debugLoadingEnabled] = useState<boolean>(() =>
112+
isStockfishDebugLoadingEnabled(),
113+
)
114+
const [loadingLabel, setLoadingLabel] = useState<string>(() =>
115+
getStockfishLoadingLabel(engineRef.current, debugLoadingEnabled),
116+
)
117+
const toastId = useRef<string | null>(null)
118+
const loadingToastTimerRef = useRef<number | null>(null)
42119

43120
const streamEvaluations = useCallback(
44121
(fen: string, legalMoveCount: number, depth?: number) => {
@@ -61,10 +138,21 @@ export const StockfishEngineContextProvider: React.FC<{
61138

62139
useEffect(() => {
63140
const checkEngineStatus = () => {
64-
if (engineRef.current?.ready) {
141+
const engine = engineRef.current
142+
if (!engine) return
143+
144+
setLoadingLabel((prev) => {
145+
const next = getStockfishLoadingLabel(engine, debugLoadingEnabled)
146+
return prev === next ? prev : next
147+
})
148+
149+
if (engine.initializationError) {
150+
setStatus('error')
151+
setError(engine.initializationError)
152+
} else if (engine.ready) {
65153
setStatus('ready')
66154
setError(null)
67-
} else if (engineRef.current && !engineRef.current.ready) {
155+
} else {
68156
setStatus('loading')
69157
}
70158
}
@@ -73,38 +161,81 @@ export const StockfishEngineContextProvider: React.FC<{
73161
const interval = setInterval(checkEngineStatus, 100)
74162

75163
return () => clearInterval(interval)
76-
}, [])
164+
}, [debugLoadingEnabled])
77165

78166
// Toast notifications for Stockfish engine status
79167
useEffect(() => {
80168
return () => {
81-
toast.dismiss()
169+
if (
170+
loadingToastTimerRef.current !== null &&
171+
typeof window !== 'undefined'
172+
) {
173+
window.clearTimeout(loadingToastTimerRef.current)
174+
loadingToastTimerRef.current = null
175+
}
176+
if (toastId.current) {
177+
toast.dismiss(toastId.current)
178+
toastId.current = null
179+
}
82180
}
83181
}, [])
84182

85183
useEffect(() => {
86-
if (status === 'loading' && !toastId.current) {
87-
toastId.current = toast.loading('Loading Stockfish Engine...')
88-
} else if (status === 'ready') {
184+
if (status === 'loading') {
185+
if (typeof window === 'undefined') {
186+
return
187+
}
188+
189+
if (toastId.current) {
190+
toastId.current = toast.loading(loadingLabel, { id: toastId.current })
191+
return
192+
}
193+
194+
if (loadingToastTimerRef.current !== null) {
195+
return
196+
}
197+
198+
loadingToastTimerRef.current = window.setTimeout(() => {
199+
loadingToastTimerRef.current = null
200+
if (!toastId.current && engineRef.current && !engineRef.current.ready) {
201+
toastId.current = toast.loading(
202+
getStockfishLoadingLabel(engineRef.current, debugLoadingEnabled),
203+
)
204+
}
205+
}, STOCKFISH_LOADING_TOAST_DELAY_MS)
206+
return
207+
}
208+
209+
if (
210+
loadingToastTimerRef.current !== null &&
211+
typeof window !== 'undefined'
212+
) {
213+
window.clearTimeout(loadingToastTimerRef.current)
214+
loadingToastTimerRef.current = null
215+
}
216+
217+
if (status === 'ready') {
218+
// Only show a success toast when we previously showed a loading toast.
89219
if (toastId.current) {
90220
toast.success('Loaded Stockfish! Engine is ready', {
91221
id: toastId.current,
92222
})
93223
toastId.current = null
94-
} else {
95-
toast.success('Loaded Stockfish! Engine is ready')
96224
}
97225
} else if (status === 'error') {
226+
const message = error
227+
? `Failed to load Stockfish engine: ${error}`
228+
: 'Failed to load Stockfish engine'
98229
if (toastId.current) {
99-
toast.error('Failed to load Stockfish engine', {
230+
toast.error(message, {
100231
id: toastId.current,
101232
})
102233
toastId.current = null
103234
} else {
104-
toast.error('Failed to load Stockfish engine')
235+
toast.error(message)
105236
}
106237
}
107-
}, [status])
238+
}, [status, error, loadingLabel, debugLoadingEnabled])
108239

109240
const contextValue = useMemo(
110241
() => ({

‎src/lib/engine/maia.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,11 @@ class Maia {
9797
position += chunk.length
9898
}
9999

100-
await this.storage.storeModel(this.modelUrl, this.modelVersion, buffer.buffer)
100+
await this.storage.storeModel(
101+
this.modelUrl,
102+
this.modelVersion,
103+
buffer.buffer,
104+
)
101105

102106
await this.initializeModel(buffer.buffer)
103107
this.options.setStatus('ready')

0 commit comments

Comments
 (0)