fix(speech): harden resumed model downloads

2026-10-18 00:03:20 +00:00 · 2026-07-14 15:27:08 -07:00
parent 26b7bdf3af
commit 9a73173d26
2 changed files with 267 additions and 30 deletions
@@ -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<void>
+  downloadFile: (
+    url: string,
+    archivePath: string,
+    expectedSize: number,
+    modelId: string,
+    isAborted: () => boolean,
+    signal?: AbortSignal,
+    redirectCount?: number,
+    resumeOffset?: number
+  ) => Promise<void>
 }

 type ScriptedResponse = {
@@ -41,7 +51,10 @@ type ScriptedResponseFactory = (sentHeaders: Record<string, string>) => 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<string, string> } {
+function scriptRequest(factory: ScriptedResponseFactory): {
+  sentHeaders: Record<string, string>
+  abortMock: ReturnType<typeof vi.fn>
+} {
  const sentHeaders: Record<string, string> = {}
  const listeners = new Map<string, Set<(...args: unknown[]) => 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

@@ -35,11 +35,15 @@ type ProgressCallback = (modelId: string, progress: number) => void
 type DownloadIncomingMessage = Electron.IncomingMessage &
  NodeJS.ReadableStream & {
    headers: Record<string, string | string[] | undefined>
-    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
        }