diff --git a/src/chrome/src/agent/apocalypse-mode.js b/src/chrome/src/agent/apocalypse-mode.js index 002bf486b..5d21af904 100644 --- a/src/chrome/src/agent/apocalypse-mode.js +++ b/src/chrome/src/agent/apocalypse-mode.js @@ -821,6 +821,7 @@ export function createOpfsArchiveStorage(storageManager = globalThis.navigator?. } const MAX_RETRY_ATTEMPTS = 6; +const DEFAULT_MAX_PIECES_PER_WAKE = 8; const BASE_RETRY_MS = 60_000; const MAX_RETRY_MS = 6 * 60 * 60_000; export const APOCALYPSE_DOWNLOAD_ALARM = 'wb_apocalypse_archive_download'; @@ -896,7 +897,9 @@ export function createApocalypseArchiveManager(options = {}) { const randomId = options.randomId || (() => globalThis.crypto.randomUUID()); const now = options.now || (() => Date.now()); const configuredMaxPieces = Number(options.maxPiecesPerWake); - const maxPiecesPerWake = Number.isFinite(configuredMaxPieces) ? Math.max(1, Math.floor(configuredMaxPieces)) : Number.POSITIVE_INFINITY; + const maxPiecesPerWake = Number.isFinite(configuredMaxPieces) + ? Math.max(1, Math.floor(configuredMaxPieces)) + : DEFAULT_MAX_PIECES_PER_WAKE; const controllers = new Map(); let processing = false; if (!store || !storage) throw new Error('Apocalypse Mode requires state and archive storage adapters.'); diff --git a/src/chrome/src/agent/emergency-box.js b/src/chrome/src/agent/emergency-box.js index 4e7d69817..31548bae8 100644 --- a/src/chrome/src/agent/emergency-box.js +++ b/src/chrome/src/agent/emergency-box.js @@ -378,7 +378,17 @@ export async function downloadEmergencyResource(resource, options = {}) { error: '', }); } catch (error) { - try { await writer?.abort?.(error); } catch { /* keep the last committed partial file */ } + if (writer) { + try { + // OPFS writable streams are atomic: aborting rolls the file back to its + // pre-download size. Commit successfully written chunks so a pause or + // transient network failure can resume from durable progress. + await writer.close(); + } catch { + try { await writer.abort?.(error); } catch { /* preserve the original failure */ } + } + writer = null; + } const bytesReceived = await storage.size(storageKey).catch(() => 0); const paused = error?.name === 'AbortError' || signal?.aborted; await persist({ diff --git a/src/chrome/src/offscreen/inference-worker.js b/src/chrome/src/offscreen/inference-worker.js index 3897f16a3..d9fd60526 100644 --- a/src/chrome/src/offscreen/inference-worker.js +++ b/src/chrome/src/offscreen/inference-worker.js @@ -957,12 +957,15 @@ self.addEventListener('message', async event => { const request = assertTextDownloadCanStart(payload); queuedTextDownload = request; let acknowledged = false; - const operation = enqueueModelOperation(() => downloadTextModel(payload, { - onStarted(state) { - acknowledged = true; - self.postMessage({ id, ok: true, ...state }); - }, - })); + const operation = enqueueModelOperation(() => { + if (queuedTextDownload !== request) return getTextDownloadStatus(request.modelId, request.dtype); + return downloadTextModel(payload, { + onStarted(state) { + acknowledged = true; + self.postMessage({ id, ok: true, ...state }); + }, + }); + }); void operation.then((state) => { if (!acknowledged) self.postMessage({ id, ok: true, ...state }); }).catch((error) => { @@ -979,6 +982,9 @@ self.addEventListener('message', async event => { if (type === 'stop-text-download') { const modelId = String(payload?.modelId || '').trim(); const dtype = payload?.dtype || 'q4f16'; + const targetsQueuedTransfer = queuedTextDownload + && sameTextModel(queuedTextDownload.modelId, queuedTextDownload.dtype, modelId, dtype); + if (targetsQueuedTransfer) queuedTextDownload = null; const targetsTrackedTransfer = sameTextModel(textDownloadState.modelId, textDownloadState.dtype, modelId, dtype); if (targetsTrackedTransfer) { textDownloadCancelMode = 'stop'; diff --git a/src/firefox/src/agent/apocalypse-mode.js b/src/firefox/src/agent/apocalypse-mode.js index 002bf486b..5d21af904 100644 --- a/src/firefox/src/agent/apocalypse-mode.js +++ b/src/firefox/src/agent/apocalypse-mode.js @@ -821,6 +821,7 @@ export function createOpfsArchiveStorage(storageManager = globalThis.navigator?. } const MAX_RETRY_ATTEMPTS = 6; +const DEFAULT_MAX_PIECES_PER_WAKE = 8; const BASE_RETRY_MS = 60_000; const MAX_RETRY_MS = 6 * 60 * 60_000; export const APOCALYPSE_DOWNLOAD_ALARM = 'wb_apocalypse_archive_download'; @@ -896,7 +897,9 @@ export function createApocalypseArchiveManager(options = {}) { const randomId = options.randomId || (() => globalThis.crypto.randomUUID()); const now = options.now || (() => Date.now()); const configuredMaxPieces = Number(options.maxPiecesPerWake); - const maxPiecesPerWake = Number.isFinite(configuredMaxPieces) ? Math.max(1, Math.floor(configuredMaxPieces)) : Number.POSITIVE_INFINITY; + const maxPiecesPerWake = Number.isFinite(configuredMaxPieces) + ? Math.max(1, Math.floor(configuredMaxPieces)) + : DEFAULT_MAX_PIECES_PER_WAKE; const controllers = new Map(); let processing = false; if (!store || !storage) throw new Error('Apocalypse Mode requires state and archive storage adapters.'); diff --git a/src/firefox/src/agent/emergency-box.js b/src/firefox/src/agent/emergency-box.js index 4e7d69817..31548bae8 100644 --- a/src/firefox/src/agent/emergency-box.js +++ b/src/firefox/src/agent/emergency-box.js @@ -378,7 +378,17 @@ export async function downloadEmergencyResource(resource, options = {}) { error: '', }); } catch (error) { - try { await writer?.abort?.(error); } catch { /* keep the last committed partial file */ } + if (writer) { + try { + // OPFS writable streams are atomic: aborting rolls the file back to its + // pre-download size. Commit successfully written chunks so a pause or + // transient network failure can resume from durable progress. + await writer.close(); + } catch { + try { await writer.abort?.(error); } catch { /* preserve the original failure */ } + } + writer = null; + } const bytesReceived = await storage.size(storageKey).catch(() => 0); const paused = error?.name === 'AbortError' || signal?.aborted; await persist({ diff --git a/test/run.js b/test/run.js index 1ed1536f9..237927e54 100644 --- a/test/run.js +++ b/test/run.js @@ -21112,6 +21112,57 @@ test('Emergency Box streams PDFs to resumable local storage and rejects non-PDF } }); +test('Emergency Box commits partial PDF bytes before recording a paused download', async () => { + for (const [label, runtime] of [['chrome', EmergencyBoxCh], ['firefox', EmergencyBoxFx]]) { + let committed = new Uint8Array(); + const records = new Map(); + const controller = new AbortController(); + const chunks = [new TextEncoder().encode('%PDF-partial'), new TextEncoder().encode('-ignored')]; + const storage = { + async size() { return committed.byteLength; }, + async createWriter() { + let working = committed.slice(); + return { + async write(position, bytes) { + const expanded = new Uint8Array(position + bytes.byteLength); + expanded.set(working); + expanded.set(bytes, position); + working = expanded; + }, + async truncate(size) { working = working.slice(0, size); }, + async close() { committed = working; }, + async abort() {}, + }; + }, + }; + let readIndex = 0; + const result = await runtime.downloadEmergencyResource({ + id: 'paused-pdf', title: 'Paused PDF', url: 'https://example.test/paused.pdf', + }, { + signal: controller.signal, + storage, + store: { + async get(id) { return records.get(id); }, + async put(record) { records.set(record.id, { ...record }); return record; }, + }, + fetchImpl: async () => ({ + ok: true, + status: 200, + headers: { get() { return null; } }, + body: { getReader() { return { async read() { + if (readIndex === 1) controller.abort(); + if (readIndex >= chunks.length) return { done: true }; + return { done: false, value: chunks[readIndex++] }; + } }; } }, + }), + }); + + assert.equal(result.status, 'paused', `${label}: aborted PDF did not become paused`); + assert.equal(result.bytesReceived, chunks[0].byteLength, `${label}: paused byte count was rolled back`); + assert.equal(committed.byteLength, chunks[0].byteLength, `${label}: partial PDF transaction was not committed`); + } +}); + test('Emergency Box UI and PDF reader stay in Chrome and Firefox parity', () => { const files = [ 'src/agent/emergency-box.js', @@ -22373,6 +22424,36 @@ test('Apocalypse Mode rolls back an interrupted OPFS write session before resumi } }); +test('Apocalypse Mode bounds production OPFS write sessions', async () => { + for (const [label, runtime] of [['chrome', ApocalypseModeCh], ['firefox', ApocalypseModeFx]]) { + const record = { + id: 'bounded-session', status: 'queued', generation: 1, updatedAt: 1, + filename: 'archive.zim', size: 10, pieceLength: 1, pieceHashAlgorithm: 'sha-1', + pieceHashes: Array(10).fill('valid'), downloadUrl: 'https://example.test/archive.zim', + target: { kind: 'opfs', key: 'archive.zim' }, pieceIndex: 0, bytesDownloaded: 0, retryCount: 0, + }; + const records = new Map([[record.id, record]]); + let closes = 0; + const manager = runtime.createApocalypseArchiveManager({ + store: { + async getConfig() { return { enabled: true }; }, + async listArchives() { return [...records.values()].map(value => ({ ...value })); }, + async getArchive(id) { const value = records.get(id); return value ? { ...value } : null; }, + async putArchive(value) { records.set(value.id, { ...value }); return value; }, + }, + storage: { async createWriter() { return { async write() {}, async close() { closes += 1; }, async abort() {} }; } }, + fetchImpl: async () => ({ ok: true, status: 206, async arrayBuffer() { return Uint8Array.of(1).buffer; } }), + digestHex: async () => 'valid', schedule() {}, randomId: () => 'bounded-lease', now: () => 100, + }); + + const result = await manager.processNext(); + + assert.equal(result.archive.status, 'queued', `${label}: default wake consumed the entire archive`); + assert.equal(result.archive.pieceIndex, 8, `${label}: default write batch was not bounded`); + assert.equal(closes, 1, `${label}: bounded write batch was not committed`); + } +}); + test('Apocalypse Mode rejects a corrupt piece before storage and backs off', async () => { for (const [label, runtime] of [['chrome', ApocalypseModeCh], ['firefox', ApocalypseModeFx]]) { const config = { enabled: true }; @@ -43850,6 +43931,8 @@ test('WebGPU worker replays text tool history and applies model-specific generat const previousReleaseTextDownload = globalThis.__releaseWebgpuTextDownload; const previousGenerationOptions = globalThis.__webgpuGenerationOptions; const previousPipelineOptions = globalThis.__webgpuPipelineOptions; + const previousHoldTextGeneration = globalThis.__holdWebgpuTextGeneration; + const previousReleaseTextGeneration = globalThis.__releaseWebgpuTextGeneration; let workerListener = null; const posted = []; try { @@ -43947,6 +44030,9 @@ test('WebGPU worker replays text tool history and applies model-specific generat await new Promise(resolve => { globalThis.__releaseWebgpuTextDownload = resolve; }); } const instance = async (input, options) => { + if (globalThis.__holdWebgpuTextGeneration) { + await new Promise(resolve => { globalThis.__releaseWebgpuTextGeneration = resolve; }); + } globalThis.__webgpuGenerationOptions = options; const content = modelId === 'LiquidAI/LFM2.5-2.6B-ONNX' ? 'private model reasoningHello!' @@ -44100,6 +44186,24 @@ test('WebGPU worker replays text tool history and applies model-specific generat assert.equal(activeStatus.status, 'ready'); assert.equal(activeStatus.ready, true); + globalThis.__holdWebgpuTextGeneration = true; + const generationId = requestId++; + const generationPromise = workerListener({ data: { id: generationId, type: 'text-chat', payload: activePayload } }); + for (let attempt = 0; attempt < 20 && !globalThis.__releaseWebgpuTextGeneration; attempt++) { + await new Promise(resolve => setTimeout(resolve, 0)); + } + const queuedPayload = { ...textPayload, modelId: 'text-model-queued' }; + const queuedId = requestId++; + const queuedPromise = workerListener({ data: { id: queuedId, type: 'start-download-text', payload: queuedPayload } }); + const stopQueuedId = requestId++; + const stopQueuedPromise = workerListener({ data: { id: stopQueuedId, type: 'stop-text-download', payload: queuedPayload } }); + globalThis.__releaseWebgpuTextGeneration(); + await Promise.all([generationPromise, queuedPromise, stopQueuedPromise]); + assert.equal(posted.find(message => message.id === queuedId)?.status, 'not-downloaded', 'stopped queued download still started'); + assert.equal(posted.find(message => message.id === stopQueuedId)?.status, 'not-downloaded', 'queued Stop did not clear the target model'); + globalThis.__holdWebgpuTextGeneration = false; + globalThis.__releaseWebgpuTextGeneration = null; + globalThis.__holdWebgpuTextDownload = false; globalThis.__releaseWebgpuTextDownload = null; const objectDtypePayload = { @@ -44163,6 +44267,10 @@ test('WebGPU worker replays text tool history and applies model-specific generat else globalThis.__webgpuGenerationOptions = previousGenerationOptions; if (previousPipelineOptions === undefined) delete globalThis.__webgpuPipelineOptions; else globalThis.__webgpuPipelineOptions = previousPipelineOptions; + if (previousHoldTextGeneration === undefined) delete globalThis.__holdWebgpuTextGeneration; + else globalThis.__holdWebgpuTextGeneration = previousHoldTextGeneration; + if (previousReleaseTextGeneration === undefined) delete globalThis.__releaseWebgpuTextGeneration; + else globalThis.__releaseWebgpuTextGeneration = previousReleaseTextGeneration; } });