From d2befaafa97c7dafd2b06093a1cb93a1a5199ba9 Mon Sep 17 00:00:00 2001 From: alichherawalla Date: Sun, 26 Jul 2026 12:39:24 +0530 Subject: [PATCH] fix(B15/G7): gate the reconcile re-unzip on extraction completeness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit reconcileFinishedImageDownloads re-unzips a `_zip_name` remnant (a mid-unzip kill) and registered the recovered model on isValidZip alone — size>0 + PK header — SKIPPING the completeness check the live/resume download paths run. A truncated-but-PK zip yields a PARTIAL mnn/qnn tree (missing base files / weight pairs) registered as usable → native crash at generation. Fix: after the recovery unzip, run the same ensureImageExtractionComplete gate; on a still-incomplete extraction, clean up (dir + zip) and do NOT mark `_ready` / register, so it resurfaces as re-downloadable. coreml is exempt (its integrity is validated in coreMLModelUtils). The recovery block is extracted into recoverImageModelFromZipRemnant + buildRecoveredImageModel — dedupes the 3× model-construction and keeps reconcile's complexity within the gate. Test drives the REAL reconcile + REAL completeness gate (only RNFS + native unzip are faked, via a path-addressed in-memory FS — order-independent, not brittle sequential mocks): a partial mnn tree is neither registered nor left on disk (red before the fix — addImageModel WAS called; verified by reverting), a complete tree still registers, and the two legacy re-unzip mechanics tests are pinned to a completeness-exempt coreml model. 210 modelManager tests pass; tsc/depcruise clean; no new lint. --- __tests__/unit/services/modelManager.test.ts | 28 ++-- .../reconcilePartialUnzipG7.test.ts | 133 ++++++++++++++++++ .../modelManager/scan.branches.test.ts | 11 +- src/services/modelManager/scan.ts | 110 ++++++++------- 4 files changed, 218 insertions(+), 64 deletions(-) create mode 100644 __tests__/unit/services/modelManager/reconcilePartialUnzipG7.test.ts diff --git a/__tests__/unit/services/modelManager.test.ts b/__tests__/unit/services/modelManager.test.ts index 3cf81be69..0575c065e 100644 --- a/__tests__/unit/services/modelManager.test.ts +++ b/__tests__/unit/services/modelManager.test.ts @@ -2649,11 +2649,21 @@ describe('ModelManager', () => { expect(result).toHaveLength(0); }); - it('re-unzips from valid zip when _zip_name present and zip valid', async () => { + // coreml-named dir → the mnn/qnn completeness gate (ensureImageExtractionComplete) is a no-op + // (coreml integrity is validated in coreMLModelUtils). Asserts the re-unzip→register RECOVERY + // mechanics; the mnn completeness gate on the re-unzip path is covered by + // reconcilePartialUnzipG7.test.ts (B15/G7). + it('re-unzips from valid zip when _zip_name present and zip valid (coreml — completeness-exempt)', async () => { setupInitialized(); + // resetAllMocks (beforeEach) clears the module mock impl — re-establish it, like the coreml + // resolve test below. + (mockedResolveCoreML as jest.Mock).mockImplementation((dir: string) => + Promise.resolve(`${dir}/model.mlpackage`), + ); mockedAsyncStorage.getItem - .mockResolvedValueOnce('[]') - .mockResolvedValueOnce('[]'); + .mockResolvedValueOnce('[]') // reconciler read + .mockResolvedValueOnce('[]') // addDownloadedImageModel read + .mockResolvedValueOnce(null); // addDownloadedImageModel save mockedRNFS.exists .mockResolvedValueOnce(true) // imageModelsDir scan @@ -2663,13 +2673,13 @@ describe('ModelManager', () => { mockedRNFS.readDir .mockResolvedValueOnce([ - { name: 'sd_v15_mnn', path: `${IMAGE_MODELS_DIR}/sd_v15_mnn`, size: 0, isFile: () => false, isDirectory: () => true }, + { name: 'sd_v15_coreml', path: `${IMAGE_MODELS_DIR}/sd_v15_coreml`, size: 0, isFile: () => false, isDirectory: () => true }, ] as any) .mockResolvedValueOnce([ - { name: 'unet.mnn', path: `${IMAGE_MODELS_DIR}/sd_v15_mnn/unet.mnn`, size: 500000000, isFile: () => true, isDirectory: () => false }, + { name: 'model.mlpackage', path: `${IMAGE_MODELS_DIR}/sd_v15_coreml/model.mlpackage`, size: 500000000, isFile: () => true, isDirectory: () => false }, ] as any); - mockedRNFS.readFile = jest.fn().mockResolvedValueOnce('sd_v15_mnn.zip'); + mockedRNFS.readFile = jest.fn().mockResolvedValueOnce('sd_v15_coreml.zip'); mockedRNFS.stat = jest.fn().mockResolvedValue({ size: 400000000, isFile: () => true } as any); mockedRNFS.read = jest.fn().mockResolvedValue('PK\x03\x04'); mockedRNFS.writeFile = jest.fn().mockResolvedValue(undefined as any); @@ -2678,11 +2688,11 @@ describe('ModelManager', () => { const result = await modelManager.reconcileFinishedImageDownloads(new Set()); expect(mockedUnzip).toHaveBeenCalledWith( - `${IMAGE_MODELS_DIR}/sd_v15_mnn.zip`, - `${IMAGE_MODELS_DIR}/sd_v15_mnn`, + `${IMAGE_MODELS_DIR}/sd_v15_coreml.zip`, + `${IMAGE_MODELS_DIR}/sd_v15_coreml`, ); expect(result).toHaveLength(1); - expect(result[0].id).toBe('sd_v15_mnn'); + expect(result[0].id).toBe('sd_v15_coreml'); }); it('deletes partial dir when _zip_name present but zip is missing', async () => { diff --git a/__tests__/unit/services/modelManager/reconcilePartialUnzipG7.test.ts b/__tests__/unit/services/modelManager/reconcilePartialUnzipG7.test.ts new file mode 100644 index 000000000..2b76d0ba7 --- /dev/null +++ b/__tests__/unit/services/modelManager/reconcilePartialUnzipG7.test.ts @@ -0,0 +1,133 @@ +/** + * B15 / G7 — reconcileFinishedImageDownloads must NOT register a partially-extracted image model. + * + * The recovery path re-unzips a `_zip_name` remnant (a mid-unzip kill) and used to register the + * model on isValidZip alone (size>0 + PK header) — skipping the completeness gate the live/resume + * paths run. A truncated-but-PK zip yields a PARTIAL mnn tree (missing base files / weight pairs), + * registered as usable → native crash at generation. The fix runs ensureImageExtractionComplete + * after the recovery unzip; a still-incomplete extraction is cleaned up and NEVER marked _ready. + * + * This drives the REAL reconcile + the REAL completeness gate (ensureImageExtractionComplete / + * validateImageModelDir are NOT mocked). Only the third-party boundary — react-native-fs and the + * native unzip — is faked, via a path-addressed in-memory FS (order-independent, not brittle + * sequential mocks). addImageModel is the injected sink, so we assert the real outcome: a partial + * model is neither registered nor left on disk; a complete one is. + */ +jest.mock('react-native-fs', () => { + const files = new Map(); + const norm = (p: string) => p.replace(/\/+$/, ''); + return { + __files: files, + exists: jest.fn(async (p: string) => files.has(norm(p))), + stat: jest.fn(async (p: string) => { + const f = files.get(norm(p)); + if (!f) throw new Error('ENOENT'); + return { size: f.size, isFile: () => f.isFile, isDirectory: () => !f.isFile }; + }), + read: jest.fn(async (p: string) => (norm(p).endsWith('.zip') ? 'PK' : '')), + readFile: jest.fn(async (p: string) => { + const f = files.get(norm(p)); + return f && (f as any).content !== undefined ? (f as any).content : ''; + }), + readDir: jest.fn(async (p: string) => { + const base = norm(p); + const seen = new Set(); + const out: any[] = []; + for (const path of files.keys()) { + if (path === base || !path.startsWith(`${base}/`)) continue; + const rest = path.slice(base.length + 1); + const first = rest.split('/')[0]; + const childPath = `${base}/${first}`; + if (seen.has(childPath)) continue; + seen.add(childPath); + const childMeta = files.get(childPath); + const isFile = rest.includes('/') ? false : (childMeta?.isFile ?? true); + out.push({ + name: first, + path: childPath, + size: childMeta?.size ?? 0, + isFile: () => isFile, + isDirectory: () => !isFile, + }); + } + return out; + }), + writeFile: jest.fn(async (p: string, content = '') => { files.set(norm(p), { isFile: true, size: content.length || 1, ...(content ? { content } : {}) } as any); }), + unlink: jest.fn(async (p: string) => { + const base = norm(p); + for (const key of [...files.keys()]) if (key === base || key.startsWith(`${base}/`)) files.delete(key); + }), + }; +}); +jest.mock('react-native-zip-archive', () => ({ unzip: jest.fn(async () => '') })); +jest.mock('../../../../src/services/modelManager/storage', () => ({ + buildDownloadedModel: jest.fn(), persistDownloadedModel: jest.fn(), + loadDownloadedModels: jest.fn(async () => []), saveModelsList: jest.fn(), +})); +jest.mock('../../../../src/utils/coreMLModelUtils', () => ({ resolveCoreMLModelDir: jest.fn(async (d: string) => d) })); + +import RNFS from 'react-native-fs'; +import { reconcileFinishedImageDownloads } from '../../../../src/services/modelManager/scan'; + +const files: Map = (RNFS as any).__files; +const IMG = '/img'; +const ID = 'Sd_Test_Mnn'; // detectBackend → 'mnn' +const DIR = `${IMG}/${ID}`; +const ZIP = `${IMG}/archive.zip`; + +const seed = (p: string, size = 1024, content?: string) => + files.set(p, { isFile: true, size, ...(content !== undefined ? { content } : {}) }); + +const COMPLETE_MNN = [ + 'pos_emb.bin', 'token_emb.bin', 'tokenizer.json', + 'unet.mnn', 'unet.mnn.weight', + 'vae_decoder.mnn', 'vae_decoder.mnn.weight', + 'clip_v2.mnn', 'clip_v2.mnn.weight', +]; + +function stageCommon() { + files.clear(); + files.set(IMG, { isFile: false, size: 0 }); + files.set(DIR, { isFile: false, size: 0 }); + seed(`${DIR}/_zip_name`, 12, 'archive.zip'); // remnant of a mid-unzip kill (no _ready) + seed(ZIP, 24 * 1024 * 1024); // valid PK zip, non-zero +} + +const opts = () => ({ + imageModelsDir: IMG, + getImageModels: jest.fn(async (): Promise => []), + addImageModel: jest.fn(async () => {}), + activeModelIds: new Set(), +}); + +describe('reconcileFinishedImageDownloads — partial re-unzip completeness gate (B15/G7)', () => { + it('does NOT register a partially-extracted mnn model, and cleans the dir up', async () => { + stageCommon(); + // Partial extraction: only the unet graph landed — base files, weight pairs, clip, vae missing. + seed(`${DIR}/unet.mnn`, 4096); + const o = opts(); + + const out = await reconcileFinishedImageDownloads(o); + + expect(o.addImageModel).not.toHaveBeenCalled(); // never published + expect(out).toEqual([]); + expect(files.has(`${DIR}/_ready`)).toBe(false); // never marked ready + expect(files.has(DIR)).toBe(false); // unusable dir removed → re-downloadable + expect(files.has(ZIP)).toBe(false); // and the bad zip cleaned up + }); + + it('DOES register a fully-extracted mnn model (gate passes — no over-block)', async () => { + stageCommon(); + COMPLETE_MNN.forEach(f => seed(`${DIR}/${f}`, 4096)); + const o = opts(); + + const out = await reconcileFinishedImageDownloads(o); + + expect(o.addImageModel).toHaveBeenCalledTimes(1); + expect(out).toHaveLength(1); + expect(out[0].id).toBe(ID); + expect(out[0].backend).toBe('mnn'); + expect(files.has(`${DIR}/_ready`)).toBe(true); // marked ready + expect(files.has(DIR)).toBe(true); // kept + }); +}); diff --git a/__tests__/unit/services/modelManager/scan.branches.test.ts b/__tests__/unit/services/modelManager/scan.branches.test.ts index 9d801c647..1cf5242e9 100644 --- a/__tests__/unit/services/modelManager/scan.branches.test.ts +++ b/__tests__/unit/services/modelManager/scan.branches.test.ts @@ -258,7 +258,10 @@ describe('reconcileFinishedImageDownloads', () => { expect(out).toEqual([]); }); - it('re-unzips when _zip_name present and zip valid', async () => { + // coreml-named dir → the mnn/qnn completeness gate is a no-op (coreml integrity is validated in + // coreMLModelUtils). The mnn completeness gate on the re-unzip path is covered by + // reconcilePartialUnzipG7.test.ts (B15). This test asserts the re-unzip→register recovery mechanics. + it('re-unzips when _zip_name present and zip valid (coreml — completeness-exempt)', async () => { const opts = base(); mockedRNFS.exists .mockResolvedValueOnce(true) // dir @@ -266,13 +269,13 @@ describe('reconcileFinishedImageDownloads', () => { .mockResolvedValueOnce(true) // _zip_name present .mockResolvedValueOnce(true); // isValidZip: zip exists mockedRNFS.readDir - .mockResolvedValueOnce([dir('Zipped', '/img/Zipped')]) - .mockResolvedValueOnce([file('w', '/img/Zipped/w', 999)]); // getDirSize + .mockResolvedValueOnce([dir('coreml_Zipped', '/img/coreml_Zipped')]) + .mockResolvedValueOnce([file('w', '/img/coreml_Zipped/w', 999)]); // getDirSize (mockedRNFS.readFile as jest.Mock).mockResolvedValueOnce('archive.zip'); (mockedRNFS.stat as jest.Mock).mockResolvedValueOnce({ size: 1024 }); (mockedRNFS.read as jest.Mock).mockResolvedValueOnce('PK'); const out = await reconcileFinishedImageDownloads(opts); - expect(unzip).toHaveBeenCalledWith('/img/archive.zip', '/img/Zipped'); + expect(unzip).toHaveBeenCalledWith('/img/archive.zip', '/img/coreml_Zipped'); expect(out).toHaveLength(1); expect(out[0].size).toBe(999); }); diff --git a/src/services/modelManager/scan.ts b/src/services/modelManager/scan.ts index 28bcf3bb4..280f38279 100644 --- a/src/services/modelManager/scan.ts +++ b/src/services/modelManager/scan.ts @@ -4,6 +4,7 @@ import { DownloadedModel, LlamaDownloadedModel, LiteRTDownloadedModel, ModelFile import { buildDownloadedModel, persistDownloadedModel, loadDownloadedModels, saveModelsList } from './storage'; import { copyFileWithProgress } from './copyFile'; import { resolveCoreMLModelDir } from '../../utils/coreMLModelUtils'; +import { ensureImageExtractionComplete } from '../../utils/imageModelIntegrity'; // Single source of truth for projector detection + model↔projector matching (see src/services/mmproj.ts). import { isMMProjFile, pickMmProjForModel } from '../mmproj'; @@ -133,6 +134,57 @@ async function isValidZip(zipPath: string): Promise { return true; } +/** Build the ONNXImageModel record for a recovered on-disk dir (coreml resolves its inner model dir). */ +async function buildRecoveredImageModel( + item: { name: string; path: string }, + backend: 'mnn' | 'qnn' | 'coreml', +): Promise { + let modelPath = item.path; + if (backend === 'coreml') modelPath = await resolveCoreMLModelDir(item.path).catch(() => item.path); + const totalSize = await getDirSize(item.path); + return { + id: item.name, + name: item.name.replaceAll('_', ' '), + description: '', + modelPath, + size: totalSize, + downloadedAt: new Date().toISOString(), + backend, + }; +} + +/** + * Recover an image model from a `_zip_name` remnant (a mid-unzip kill): re-unzip and register it — + * but ONLY if the extraction is complete. isValidZip checks the PK header + size>0 only, so a + * truncated-but-PK zip yields a PARTIAL mnn/qnn tree that crashes natively at generation (G7). Run + * the same completeness gate the live/resume paths run; on a still-incomplete extraction, clean up + * and return null so the dir resurfaces as re-downloadable. Returns null (and cleans up) on any + * unrecoverable state; never marks `_ready` for a partial model. + */ +async function recoverImageModelFromZipRemnant( + item: { name: string; path: string }, + imageModelsDir: string, +): Promise { + const zipFileName = (await RNFS.readFile(`${item.path}/_zip_name`, 'utf8')).trim(); + const zipPath = `${imageModelsDir}/${zipFileName}`; + if (!(await isValidZip(zipPath))) { + await RNFS.unlink(item.path).catch(() => {}); // zip gone or corrupt — partial dir unrecoverable + return null; + } + await unzip(zipPath, item.path); + const backend = detectBackend(item.name); + try { + await ensureImageExtractionComplete({ backend, modelDir: item.path, zipPath, modelId: item.name }); + } catch { + await RNFS.unlink(item.path).catch(() => {}); + await RNFS.unlink(zipPath).catch(() => {}); + return null; + } + await RNFS.unlink(zipPath).catch(() => {}); + await RNFS.writeFile(`${item.path}/_ready`, '', 'utf8').catch(() => {}); + return buildRecoveredImageModel(item, backend); +} + export async function reconcileFinishedImageDownloads(opts: ReconcileImageModelsOpts): Promise { const { imageModelsDir, getImageModels, addImageModel, activeModelIds } = opts; const recovered: ONNXImageModel[] = []; @@ -188,63 +240,19 @@ export async function reconcileFinishedImageDownloads(opts: ReconcileImageModels if (hasReady) { // Unzip completed but registerAndNotify was killed — register now. - const backend = detectBackend(item.name); - let modelPath = item.path; - if (backend === 'coreml') { - modelPath = await resolveCoreMLModelDir(item.path).catch(() => item.path); - } - const totalSize = await getDirSize(item.path); - const newModel: ONNXImageModel = { - id: item.name, - name: item.name.replaceAll('_', ' '), - description: '', - modelPath, - size: totalSize, - downloadedAt: new Date().toISOString(), - backend, - }; + const newModel = await buildRecoveredImageModel(item, detectBackend(item.name)); await addImageModel(newModel); recovered.push(newModel); continue; } // No _ready — check if a zip exists to re-unzip (mid-unzip kill). - const zipNamePath = `${item.path}/_zip_name`; - const hasZipName = await RNFS.exists(zipNamePath); - - if (hasZipName) { - try { - const zipFileName = (await RNFS.readFile(zipNamePath, 'utf8')).trim(); - const zipPath = `${imageModelsDir}/${zipFileName}`; - const zipOk = await isValidZip(zipPath); - - if (zipOk) { - await unzip(zipPath, item.path); - await RNFS.unlink(zipPath).catch(() => {}); - await RNFS.writeFile(readyPath, '', 'utf8').catch(() => {}); - const backend = detectBackend(item.name); - let modelPath = item.path; - if (backend === 'coreml') { - modelPath = await resolveCoreMLModelDir(item.path).catch(() => item.path); - } - const totalSize = await getDirSize(item.path); - const newModel: ONNXImageModel = { - id: item.name, - name: item.name.replaceAll('_', ' '), - description: '', - modelPath, - size: totalSize, - downloadedAt: new Date().toISOString(), - backend, - }; - await addImageModel(newModel); - recovered.push(newModel); - } else { - // Zip is gone or corrupt — partial dir is unrecoverable, clean up. - await RNFS.unlink(item.path).catch(() => {}); - } - } catch { - // Non-fatal: leave for the next startup attempt. + if (await RNFS.exists(`${item.path}/_zip_name`)) { + // Non-fatal on unexpected error: leave the dir for the next startup attempt. + const model = await recoverImageModelFromZipRemnant(item, imageModelsDir).catch(() => null); + if (model) { + await addImageModel(model); + recovered.push(model); } } else { // No _ready and no _zip_name — stale artifact from a cancelled or