From 9a73173d26652b4aaddadeeedbde81c05a99e4e9 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Tue, 14 Jul 2026 15:27:08 -0700 Subject: [PATCH] fix(speech): harden resumed model downloads --- .../model-manager-download-resume.test.ts | 181 ++++++++++++++++-- src/main/speech/model-manager.ts | 116 +++++++++-- 2 files changed, 267 insertions(+), 30 deletions(-) diff --git a/src/main/speech/model-manager-download-resume.test.ts b/src/main/speech/model-manager-download-resume.test.ts index d334ac1147a..d42c1fc76b9 100644 --- a/src/main/speech/model-manager-download-resume.test.ts +++ b/src/main/speech/model-manager-download-resume.test.ts @@ -1,8 +1,8 @@ -import { mkdtempSync, readFileSync, rmSync } from 'node:fs' +import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from 'node:fs' import { tmpdir } from 'node:os' import { join } from 'node:path' import { Readable } from 'node:stream' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { ModelManager } from './model-manager' const { netRequestMock } = vi.hoisted(() => ({ @@ -27,6 +27,16 @@ type ModelManagerInternals = { isAborted: () => boolean, signal: AbortSignal ) => Promise + downloadFile: ( + url: string, + archivePath: string, + expectedSize: number, + modelId: string, + isAborted: () => boolean, + signal?: AbortSignal, + redirectCount?: number, + resumeOffset?: number + ) => Promise } type ScriptedResponse = { @@ -41,7 +51,10 @@ type ScriptedResponseFactory = (sentHeaders: Record) => Scripted // Emulates Electron's ClientRequest/IncomingMessage closely enough for the // download pipeline: a real Readable body so stream.pipeline semantics // (including mid-body destroy) match production behavior. -function scriptRequest(factory: ScriptedResponseFactory): { sentHeaders: Record } { +function scriptRequest(factory: ScriptedResponseFactory): { + sentHeaders: Record + abortMock: ReturnType +} { const sentHeaders: Record = {} const listeners = new Map void>>() const addListener = (event: string, cb: (...args: unknown[]) => void): void => { @@ -49,6 +62,7 @@ function scriptRequest(factory: ScriptedResponseFactory): { sentHeaders: Record< set.add(cb) listeners.set(event, set) } + const abortMock = vi.fn(() => request) const request = { setHeader: vi.fn((name: string, value: string) => { sentHeaders[name.toLowerCase()] = value @@ -61,7 +75,7 @@ function scriptRequest(factory: ScriptedResponseFactory): { sentHeaders: Record< listeners.get(event)?.delete(cb) return request }), - abort: vi.fn(() => request), + abort: abortMock, end: vi.fn(() => { queueMicrotask(() => { const spec = factory(sentHeaders) @@ -91,7 +105,7 @@ function scriptRequest(factory: ScriptedResponseFactory): { sentHeaders: Record< }) } netRequestMock.mockImplementationOnce(() => request) - return { sentHeaders } + return { sentHeaders, abortMock } } const PAYLOAD = Buffer.from('0123456789abcdefghij') @@ -101,6 +115,11 @@ describe('ModelManager download resume', () => { netRequestMock.mockReset() }) + afterEach(() => { + vi.useRealTimers() + vi.restoreAllMocks() + }) + it('resumes an interrupted download with a Range request and assembles the full file', async () => { const dir = mkdtempSync(join(tmpdir(), 'orca-model-resume-')) try { @@ -114,7 +133,10 @@ describe('ModelManager download resume', () => { expect(sentHeaders.range).toBe('bytes=10-') return { statusCode: 206, - headers: { 'content-length': String(PAYLOAD.length - 10) }, + headers: { + 'content-length': String(PAYLOAD.length - 10), + 'content-range': `bytes 10-${PAYLOAD.length - 1}/${PAYLOAD.length}` + }, chunks: [PAYLOAD.subarray(10)] } }) @@ -138,6 +160,45 @@ describe('ModelManager download resume', () => { } }) + it('rejects and discards a partial when Content-Range does not match the offset', async () => { + const dir = mkdtempSync(join(tmpdir(), 'orca-model-resume-')) + try { + const mismatched = scriptRequest(() => ({ + statusCode: 206, + headers: { + 'content-length': '10', + 'content-range': `bytes 0-9/${PAYLOAD.length}` + }, + chunks: [PAYLOAD.subarray(0, 10)] + })) + const manager = new ModelManager(dir) as unknown as ModelManagerInternals + const archivePath = join(dir, 'model.tar.bz2') + writeFileSync(archivePath, PAYLOAD.subarray(0, 10)) + + const error = await manager + .downloadFile( + 'https://example.com/model.tar.bz2', + archivePath, + PAYLOAD.length, + 'm', + () => false, + new AbortController().signal, + 0, + 10 + ) + .catch((cause: unknown) => cause) + + expect(error).toMatchObject({ + message: 'Invalid Content-Range for resume at byte 10', + retryable: true + }) + expect(existsSync(archivePath)).toBe(false) + expect(mismatched.abortMock).toHaveBeenCalledTimes(1) + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + it('restarts from scratch when the server ignores the Range request', async () => { const dir = mkdtempSync(join(tmpdir(), 'orca-model-resume-')) try { @@ -194,18 +255,100 @@ describe('ModelManager download resume', () => { } }) + it('aborts an HTTP error response instead of draining it after rejection', async () => { + const dir = mkdtempSync(join(tmpdir(), 'orca-model-resume-')) + try { + const scripted = scriptRequest(() => ({ + statusCode: 429, + headers: { 'retry-after': '3' } + })) + const manager = new ModelManager(dir) as unknown as ModelManagerInternals + + const error = await manager + .downloadFile( + 'https://example.com/model.tar.bz2', + join(dir, 'model.tar.bz2'), + PAYLOAD.length, + 'm', + () => false + ) + .catch((cause: unknown) => cause) + + expect(error).toMatchObject({ message: 'HTTP 429', httpStatusCode: 429, retryAfterMs: 3_000 }) + expect(scripted.abortMock).toHaveBeenCalledTimes(1) + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + + it('honors Retry-After before issuing another request', async () => { + vi.useFakeTimers() + const dir = mkdtempSync(join(tmpdir(), 'orca-model-resume-')) + try { + const manager = new ModelManager(dir) as unknown as ModelManagerInternals + const rateLimitError = Object.assign(new Error('HTTP 429'), { + httpStatusCode: 429, + retryAfterMs: 3_000 + }) + const downloadFileMock = vi + .spyOn(manager, 'downloadFile') + .mockRejectedValueOnce(rateLimitError) + .mockResolvedValueOnce() + const download = manager.downloadArchiveWithRetry( + 'https://example.com/model.tar.bz2', + join(dir, 'model.tar.bz2'), + PAYLOAD.length, + 'm', + () => false, + new AbortController().signal + ) + + await vi.advanceTimersByTimeAsync(2_999) + expect(downloadFileMock).toHaveBeenCalledTimes(1) + await vi.advanceTimersByTimeAsync(1) + + await expect(download).resolves.toBeUndefined() + expect(downloadFileMock).toHaveBeenCalledTimes(2) + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + + it('does not retry before an excessively long Retry-After window', async () => { + const dir = mkdtempSync(join(tmpdir(), 'orca-model-resume-')) + try { + const manager = new ModelManager(dir) as unknown as ModelManagerInternals + const rateLimitError = Object.assign(new Error('HTTP 429'), { + httpStatusCode: 429, + retryAfterMs: 300_000 + }) + const downloadFileMock = vi.spyOn(manager, 'downloadFile').mockRejectedValue(rateLimitError) + + await expect( + manager.downloadArchiveWithRetry( + 'https://example.com/model.tar.bz2', + join(dir, 'model.tar.bz2'), + PAYLOAD.length, + 'm', + () => false, + new AbortController().signal + ) + ).rejects.toThrow('server requested retry after 300 seconds') + + expect(downloadFileMock).toHaveBeenCalledTimes(1) + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + it('gives up after repeated zero-progress failures with a diagnosable error', async () => { vi.useFakeTimers() const dir = mkdtempSync(join(tmpdir(), 'orca-model-resume-')) try { - for (let i = 0; i < 8; i++) { - scriptRequest(() => ({ - statusCode: 200, - headers: { 'content-length': String(PAYLOAD.length) }, - failWith: 'net::ERR_CONNECTION_RESET' - })) - } const manager = new ModelManager(dir) as unknown as ModelManagerInternals + const downloadFileMock = vi + .spyOn(manager, 'downloadFile') + .mockRejectedValue(new Error('net::ERR_CONNECTION_RESET')) const download = manager.downloadArchiveWithRetry( 'https://example.com/model.tar.bz2', @@ -219,14 +362,16 @@ describe('ModelManager download resume', () => { () => 'resolved', (error: unknown) => (error instanceof Error ? error.message : String(error)) ) - await vi.advanceTimersByTimeAsync(60_000) + // Let the first rejection schedule its backoff before advancing timers; + // this avoids racing fake time against stream/file I/O under full-suite load. + await Promise.resolve() + await vi.runAllTimersAsync() await expect(outcome).resolves.toMatch( /Model download interrupted at 0% \(0 of 20 bytes\) after 4 attempts: .*net::ERR_CONNECTION_RESET/ ) - expect(netRequestMock).toHaveBeenCalledTimes(4) + expect(downloadFileMock).toHaveBeenCalledTimes(4) } finally { - vi.useRealTimers() rmSync(dir, { recursive: true, force: true }) } }) @@ -254,8 +399,8 @@ describe('ModelManager download resume', () => { () => 'resolved', (error: unknown) => (error instanceof Error ? error.message : String(error)) ) - // Abort while the retry loop is in its first backoff sleep; the next - // attempt must settle as Aborted without issuing another request. + // Abort after the transfer failure enters backoff; the next attempt must + // settle as Aborted without issuing another request. setTimeout(() => controller.abort(), 100) const message = await outcome diff --git a/src/main/speech/model-manager.ts b/src/main/speech/model-manager.ts index 3272a646e0c..8445330abc8 100644 --- a/src/main/speech/model-manager.ts +++ b/src/main/speech/model-manager.ts @@ -35,11 +35,15 @@ type ProgressCallback = (modelId: string, progress: number) => void type DownloadIncomingMessage = Electron.IncomingMessage & NodeJS.ReadableStream & { headers: Record - resume: () => void destroy?: () => void } -type HttpStatusError = Error & { httpStatusCode?: number } +type HttpStatusError = Error & { + httpStatusCode?: number + retryAfterMs?: number + retryable?: boolean +} type DownloadTotals = { totalBytes: number } +type ContentRange = { start: number; end: number; totalBytes?: number } const DOWNLOAD_IDLE_TIMEOUT_MS = 120_000 // Why: flaky networks and scanning proxies commonly kill long CDN transfers @@ -48,6 +52,9 @@ const DOWNLOAD_IDLE_TIMEOUT_MS = 120_000 // every full restart. const DOWNLOAD_RETRY_DELAYS_MS = [1_000, 2_000, 4_000] const MAX_DOWNLOAD_ATTEMPTS = 8 +// Why: a longer server window should be surfaced for a later manual retry, +// rather than leaving the settings download looking stuck for minutes. +const MAX_RETRY_AFTER_MS = 120_000 const RETRYABLE_NET_ERROR = /net::ERR_(CONTENT_LENGTH_MISMATCH|INCOMPLETE_CHUNKED_ENCODING|CONNECTION_(RESET|CLOSED|ABORTED|REFUSED|TIMED_OUT)|EMPTY_RESPONSE|NETWORK_CHANGED|TIMED_OUT|INTERNET_DISCONNECTED|ADDRESS_UNREACHABLE|NAME_NOT_RESOLVED|SOCKET_NOT_CONNECTED|HTTP2_PROTOCOL_ERROR|QUIC_PROTOCOL_ERROR)\b/ const RETRYABLE_HTTP_STATUSES = new Set([408, 416, 425, 429, 500, 502, 503, 504]) @@ -56,7 +63,11 @@ function isRetryableDownloadError(error: unknown): boolean { if (!(error instanceof Error)) { return false } - const statusCode = (error as HttpStatusError).httpStatusCode + const downloadError = error as HttpStatusError + if (downloadError.retryable === true) { + return true + } + const statusCode = downloadError.httpStatusCode if (statusCode !== undefined) { return RETRYABLE_HTTP_STATUSES.has(statusCode) } @@ -65,6 +76,45 @@ function isRetryableDownloadError(error: unknown): boolean { ) } +function getHeaderValue(value: string | string[] | undefined): string | undefined { + return Array.isArray(value) ? value[0] : value +} + +function parseContentRange(value: string | string[] | undefined): ContentRange | null { + const match = getHeaderValue(value) + ?.trim() + .match(/^bytes\s+(\d+)-(\d+)\/(\d+|\*)$/i) + if (!match) { + return null + } + const start = Number.parseInt(match[1], 10) + const end = Number.parseInt(match[2], 10) + const totalBytes = match[3] === '*' ? undefined : Number.parseInt(match[3], 10) + if ( + !Number.isSafeInteger(start) || + !Number.isSafeInteger(end) || + end < start || + (totalBytes !== undefined && (!Number.isSafeInteger(totalBytes) || totalBytes <= end)) + ) { + return null + } + return { start, end, totalBytes } +} + +function parseRetryAfterMs(value: string | string[] | undefined): number | undefined { + const header = getHeaderValue(value)?.trim() + if (!header) { + return undefined + } + if (/^\d+$/.test(header)) { + const seconds = Number.parseInt(header, 10) + const delayMs = seconds * 1_000 + return Number.isSafeInteger(delayMs) ? delayMs : undefined + } + const retryAt = Date.parse(header) + return Number.isNaN(retryAt) ? undefined : Math.max(0, retryAt - Date.now()) +} + function describeInterruptedDownload( cause: unknown, receivedBytes: number, @@ -436,9 +486,19 @@ export class ModelManager { ) { throw describeInterruptedDownload(err, receivedBytes, totals.totalBytes, attempt) } + const retryAfterMs = (err as HttpStatusError).retryAfterMs + if (retryAfterMs !== undefined && retryAfterMs > MAX_RETRY_AFTER_MS) { + const statusCode = (err as HttpStatusError).httpStatusCode + throw new Error( + `HTTP ${statusCode}; server requested retry after ${Math.ceil(retryAfterMs / 1_000)} seconds` + ) + } console.warn(`[speech] Model download attempt ${attempt} failed, retrying:`, modelId, err) await sleepUnlessAborted( - DOWNLOAD_RETRY_DELAYS_MS[Math.min(noProgressStreak, DOWNLOAD_RETRY_DELAYS_MS.length - 1)], + retryAfterMs ?? + DOWNLOAD_RETRY_DELAYS_MS[ + Math.min(noProgressStreak, DOWNLOAD_RETRY_DELAYS_MS.length - 1) + ], signal ) } @@ -572,9 +632,36 @@ export class ModelManager { } const onResponse = (incoming: Electron.IncomingMessage): void => { const response = incoming as DownloadIncomingMessage - const resumed = resumeOffset > 0 && response.statusCode === 206 + const contentLength = response.headers['content-length'] + const headerLength = Number.parseInt(getHeaderValue(contentLength) || '0', 10) + const parsedLength = + Number.isSafeInteger(headerLength) && headerLength > 0 ? headerLength : 0 + const contentRange = parseContentRange(response.headers['content-range']) + const resumed = + resumeOffset > 0 && + response.statusCode === 206 && + contentRange?.start === resumeOffset && + (parsedLength <= 0 || parsedLength === contentRange.end - contentRange.start + 1) + + if (resumeOffset > 0 && response.statusCode === 206 && !resumed) { + // Why: appending an unverified range can silently corrupt the archive; + // discard it so the canonical URL is retried from byte zero. + try { + rmSync(dest) + } catch { + // best-effort + } + const activeRequest = request + const rangeError: HttpStatusError = new Error( + `Invalid Content-Range for resume at byte ${resumeOffset}` + ) + rangeError.retryable = true + rejectOnce(rangeError) + activeRequest?.abort() + return + } + if (response.statusCode !== 200 && !resumed) { - response.resume() if (response.statusCode === 416) { // Why: the server rejected our resume offset; drop the partial so // the retry loop can restart this download from scratch. @@ -584,21 +671,26 @@ export class ModelManager { // best-effort } } + const activeRequest = request const statusError: HttpStatusError = new Error(`HTTP ${response.statusCode}`) statusError.httpStatusCode = response.statusCode + statusError.retryAfterMs = parseRetryAfterMs(response.headers['retry-after']) rejectOnce(statusError) + // Why: retrying must not leave a prior error body draining without + // cancellation or idle-timeout ownership. + activeRequest?.abort() return } // Why: a 200 despite our Range request means the server restarted the // transfer from byte zero, so the partial file must be overwritten. const progressBase = resumed ? resumeOffset : 0 - const contentLength = response.headers['content-length'] - const parsedLength = Number.parseInt( - Array.isArray(contentLength) ? contentLength[0] : contentLength || '0', - 10 - ) - const totalSize = parsedLength > 0 ? progressBase + parsedLength : expectedSize + const totalSize = + resumed && contentRange?.totalBytes !== undefined + ? contentRange.totalBytes + : parsedLength > 0 + ? progressBase + parsedLength + : expectedSize if (totals) { totals.totalBytes = totalSize }