Skip to content

Commit 3c2b4bb

Browse files
ashtonandersonashtonanderson
authored
Merge pull request #275 from CSSLab/codex/firefox-maia-stockfish-loading-fix
fix: unblock firefox engine loading after lichess login
2 parents 6e95c5a + 353d1d3 commit 3c2b4bb

5 files changed

Lines changed: 141 additions & 37 deletions

File tree

‎public/maia-worker.js‎

Lines changed: 35 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -116,38 +116,50 @@ self.onmessage = async (e) => {
116116
}
117117

118118
case 'download': {
119+
postMessage({ type: 'status', status: 'downloading' })
120+
postMessage({ type: 'progress', progress: 0 })
119121
const response = await fetch(modelUrl)
120122
if (!response.ok) throw new Error('Failed to fetch model')
121123

122-
const reader = response.body.getReader()
123-
const contentLength = +(response.headers.get('Content-Length') || 0)
124-
const chunks = []
125-
let receivedLength = 0
126-
let lastReportedProgress = 0
127-
128-
while (true) {
129-
const { done, value } = await reader.read()
130-
if (done) break
131-
chunks.push(value)
132-
receivedLength += value.length
133-
const currentProgress = Math.floor(
134-
(receivedLength / contentLength) * 100,
135-
)
136-
if (currentProgress >= lastReportedProgress + 10) {
137-
postMessage({ type: 'progress', progress: currentProgress })
138-
lastReportedProgress = currentProgress
124+
let buffer
125+
126+
if (response.body && typeof response.body.getReader === 'function') {
127+
const reader = response.body.getReader()
128+
const contentLength = +(response.headers.get('Content-Length') || 0)
129+
const chunks = []
130+
let receivedLength = 0
131+
let lastReportedProgress = 0
132+
133+
while (true) {
134+
const { done, value } = await reader.read()
135+
if (done) break
136+
chunks.push(value)
137+
receivedLength += value.length
138+
139+
if (contentLength > 0) {
140+
const currentProgress = Math.floor(
141+
(receivedLength / contentLength) * 100,
142+
)
143+
if (currentProgress >= lastReportedProgress + 10) {
144+
postMessage({ type: 'progress', progress: currentProgress })
145+
lastReportedProgress = currentProgress
146+
}
147+
}
139148
}
140-
}
141149

142-
const buffer = new Uint8Array(receivedLength)
143-
let position = 0
144-
for (const chunk of chunks) {
145-
buffer.set(chunk, position)
146-
position += chunk.length
150+
buffer = new Uint8Array(receivedLength)
151+
let position = 0
152+
for (const chunk of chunks) {
153+
buffer.set(chunk, position)
154+
position += chunk.length
155+
}
156+
} else {
157+
buffer = new Uint8Array(await response.arrayBuffer())
147158
}
148159

149160
await storeModel(modelUrl, modelVersion, buffer.buffer)
150161
await initSession(buffer.buffer)
162+
postMessage({ type: 'progress', progress: 100 })
151163
postMessage({ type: 'status', status: 'ready' })
152164
break
153165
}

‎src/components/Common/DownloadModelModal.tsx‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -126,14 +126,18 @@ export const DownloadModelModal: React.FC<Props> = ({
126126
</div>
127127

128128
<div className="mt-4 flex w-full flex-col items-end justify-end gap-2 md:mt-6 md:flex-row">
129-
{progress ? (
129+
{isDownloading || progress > 0 ? (
130130
<div className="relative order-2 flex h-8 w-full items-center overflow-hidden rounded-md border border-glass-border bg-glass px-3 md:order-1 md:h-10 md:flex-1">
131131
<p className="z-10 text-xs text-white/90 md:text-sm">
132-
{Math.round(progress)}%
132+
{progress > 0
133+
? `${Math.round(progress)}%`
134+
: 'Starting download...'}
133135
</p>
134136
<div
135-
className="absolute left-0 top-0 z-0 h-full rounded-l-md bg-human-4 transition-all duration-500 ease-out"
136-
style={{ width: `${progress}%` }}
137+
className={`absolute left-0 top-0 z-0 h-full rounded-l-md bg-human-4 transition-all duration-500 ease-out ${
138+
progress === 0 ? 'animate-pulse' : ''
139+
}`}
140+
style={{ width: `${progress > 0 ? progress : 12}%` }}
137141
/>
138142
</div>
139143
) : null}

‎src/lib/engine/maia.ts‎

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,18 @@ interface PendingInference {
2424
reject: (error: Error) => void
2525
}
2626

27+
interface PendingDownload {
28+
resolve: () => void
29+
reject: (error: Error) => void
30+
}
31+
2732
class Maia {
2833
private worker: Worker | null = null
2934
private options: MaiaOptions
3035
private storage: MaiaModelStorage
3136
private pendingInferences: Map<number, PendingInference> = new Map()
37+
private pendingDownload: PendingDownload | null = null
38+
private downloadPromise: Promise<void> | null = null
3239
private nextRequestId = 0
3340

3441
constructor(options: MaiaOptions) {
@@ -50,6 +57,12 @@ class Maia {
5057
switch (msg.type) {
5158
case 'status':
5259
this.options.setStatus(msg.status)
60+
if (msg.status === 'ready') {
61+
this.options.setProgress(100)
62+
this.pendingDownload?.resolve()
63+
this.pendingDownload = null
64+
this.downloadPromise = null
65+
}
5366
break
5467

5568
case 'progress':
@@ -66,6 +79,9 @@ class Maia {
6679
} else {
6780
this.options.setError(msg.message)
6881
this.options.setStatus('error')
82+
this.pendingDownload?.reject(new Error(msg.message))
83+
this.pendingDownload = null
84+
this.downloadPromise = null
6985
}
7086
break
7187
}
@@ -86,16 +102,31 @@ class Maia {
86102

87103
this.worker.onerror = (err) => {
88104
console.error('Maia worker error:', err)
89-
this.options.setError(err.message || 'Worker crashed')
105+
const error = new Error(err.message || 'Worker crashed')
106+
this.options.setError(error.message)
90107
this.options.setStatus('error')
108+
this.pendingDownload?.reject(error)
109+
this.pendingDownload = null
110+
this.downloadPromise = null
91111
}
92112

93113
this.worker.postMessage({ type: 'init', modelUrl, modelVersion })
94114
}
95115

96116
public async downloadModel() {
97117
if (!this.worker) throw new Error('Worker not initialized')
98-
this.worker.postMessage({ type: 'download' })
118+
if (this.downloadPromise) {
119+
return this.downloadPromise
120+
}
121+
122+
this.options.setProgress(0)
123+
124+
this.downloadPromise = new Promise<void>((resolve, reject) => {
125+
this.pendingDownload = { resolve, reject }
126+
this.worker!.postMessage({ type: 'download' })
127+
})
128+
129+
return this.downloadPromise
99130
}
100131

101132
public async getStorageInfo() {

‎src/lib/engine/stockfish.ts‎

Lines changed: 64 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@ import {
1111
import { StockfishModelStorage } from './stockfishStorage'
1212

1313
const DEFAULT_NNUE_FETCH_TIMEOUT_MS = 30000
14+
const DEFAULT_STOCKFISH_MODULE_INIT_TIMEOUT_MS = 15000
15+
const STOCKFISH_CACHE_LOOKUP_TIMEOUT_MS = 1500
16+
const STOCKFISH_CACHE_WRITE_TIMEOUT_MS = 5000
1417
type StockfishInitPhase =
1518
| 'idle'
1619
| 'loading-module'
@@ -1784,6 +1787,39 @@ const fetchWithTimeout = async (
17841787
}
17851788
}
17861789

1790+
const withTimeout = async <T>(
1791+
promise: Promise<T>,
1792+
timeoutMs: number,
1793+
message: string,
1794+
): Promise<T> => {
1795+
if (timeoutMs <= 0) {
1796+
return promise
1797+
}
1798+
1799+
let timeoutId: ReturnType<typeof setTimeout> | null = null
1800+
1801+
try {
1802+
return await Promise.race([
1803+
promise,
1804+
new Promise<T>((_, reject) => {
1805+
timeoutId = setTimeout(() => reject(new Error(message)), timeoutMs)
1806+
}),
1807+
])
1808+
} finally {
1809+
if (timeoutId) {
1810+
clearTimeout(timeoutId)
1811+
}
1812+
}
1813+
}
1814+
1815+
const getStockfishModuleInitTimeoutMs = (): number => {
1816+
const raw = process.env.NEXT_PUBLIC_STOCKFISH_MODULE_INIT_TIMEOUT_MS
1817+
const parsed = raw ? parseInt(raw, 10) : NaN
1818+
return Number.isFinite(parsed) && parsed > 0
1819+
? parsed
1820+
: DEFAULT_STOCKFISH_MODULE_INIT_TIMEOUT_MS
1821+
}
1822+
17871823
const loadNnueModel = async (
17881824
modelUrl: string,
17891825
storage: StockfishModelStorage,
@@ -1792,7 +1828,14 @@ const loadNnueModel = async (
17921828
forceRefresh = false,
17931829
): Promise<ArrayBuffer> => {
17941830
if (!forceRefresh) {
1795-
const cachedModel = await storage.getModel(modelUrl)
1831+
const cachedModel = await withTimeout(
1832+
storage.getModel(modelUrl),
1833+
STOCKFISH_CACHE_LOOKUP_TIMEOUT_MS,
1834+
`Timed out while checking cached Stockfish model: ${modelUrl}`,
1835+
).catch((error) => {
1836+
console.warn(error)
1837+
return null
1838+
})
17961839
if (cachedModel) {
17971840
return cachedModel
17981841
}
@@ -1807,7 +1850,13 @@ const loadNnueModel = async (
18071850
}
18081851

18091852
const buffer = await response.arrayBuffer()
1810-
await storage.storeModel(modelUrl, buffer)
1853+
void withTimeout(
1854+
storage.storeModel(modelUrl, buffer),
1855+
STOCKFISH_CACHE_WRITE_TIMEOUT_MS,
1856+
`Timed out while caching Stockfish model: ${modelUrl}`,
1857+
).catch((error) => {
1858+
console.warn(error)
1859+
})
18111860
return buffer
18121861
}
18131862

@@ -1839,16 +1888,24 @@ const setupStockfish = async (
18391888
process.env.NEXT_PUBLIC_STOCKFISH_NNUE_BASE_URL ??
18401889
'https://raw.githubusercontent.com/CSSLab/maia-platform-frontend/e23a50e/public/stockfish'
18411890
const storage = new StockfishModelStorage()
1842-
await storage.requestPersistentStorage()
18431891
const timeoutMs = getNnueFetchTimeoutMs()
1892+
const moduleInitTimeoutMs = getStockfishModuleInitTimeoutMs()
18441893
let nnueUrls: [string, string] | null = null
18451894

1895+
void storage.requestPersistentStorage().catch((error) => {
1896+
console.warn('Failed to request persistent storage for Stockfish:', error)
1897+
})
1898+
18461899
const createInstance = async (): Promise<StockfishWeb> => {
18471900
onPhaseChange?.('loading-module')
1848-
return makeModule.default({
1849-
wasmMemory: sharedWasmMemory(2560),
1850-
locateFile: (name: string) => `/stockfish/${name}`,
1851-
})
1901+
return withTimeout(
1902+
makeModule.default({
1903+
wasmMemory: sharedWasmMemory(2560),
1904+
locateFile: (name: string) => `/stockfish/${name}`,
1905+
}),
1906+
moduleInitTimeoutMs,
1907+
`Stockfish module initialization timed out after ${moduleInitTimeoutMs}ms`,
1908+
)
18521909
}
18531910

18541911
const loadWeightsIntoInstance = async (

‎src/types/engine.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ export interface MaiaEngine {
2626
maia?: Maia
2727
status: MaiaStatus
2828
progress: number
29-
downloadModel: () => void
29+
downloadModel: () => Promise<void>
3030
}
3131

3232
export interface StockfishEngine {

0 commit comments

Comments
 (0)