diff --git a/packages/mnemopi/src/core/embeddings.ts b/packages/mnemopi/src/core/embeddings.ts index 09db3e136..8e32f6551 100644 --- a/packages/mnemopi/src/core/embeddings.ts +++ b/packages/mnemopi/src/core/embeddings.ts @@ -1,4 +1,6 @@ import { mkdirSync } from "node:fs"; +import * as fsp from "node:fs/promises"; +import * as nodePath from "node:path"; import { type ApiKey, getOpenRouterHeaders, withAuth } from "@oh-my-pi/pi-ai"; import { ProviderHttpError } from "@oh-my-pi/pi-ai/error"; import { hostMatchesUrl } from "@oh-my-pi/pi-catalog/hosts"; @@ -62,12 +64,43 @@ const queryCache = new LRUCache({ max: QUERY_CACHE_MAX }); const providerIds = new WeakMap(); let nextProviderId = 1; +/** + * Quarantine the exact ONNX file named by a "Protobuf parsing failed" init + * error. A truncated cached model blocks local embeddings forever: the + * downloader treats the existing file as complete, so every init re-parses + * the same broken bytes. The extracted path is error-message CONTENT, so it + * is only honored when it resolves inside the fastembed cache directory — + * never rename an arbitrary file a dependency happens to mention. Atomic + * rename; losing a concurrent-heal race (file already renamed/removed) + * still returns true because a retry is safe either way. + * @internal exported for tests + */ +export async function quarantineCorruptModelFile(message: string, cacheDir?: string): Promise { + const match = /Load model from (.+?\.onnx) failed:.*Protobuf parsing failed/i.exec(message); + if (!match) return false; + const modelFile = nodePath.resolve(match[1]); + const cacheRoot = nodePath.resolve(cacheDir ?? getFastembedCacheDir()); + if (!modelFile.startsWith(cacheRoot + nodePath.sep)) return false; + try { + await fsp.rename(modelFile, `${modelFile}.corrupt-${Date.now()}`); + logger.warn("mnemopi: quarantined corrupt local embedding model; retrying init", { modelFile }); + } catch { + // Concurrent heal or vanished file: the single retry stays safe. A + // rename that failed with the file still in place just makes the + // retry surface the original error again. + } + return true; +} + async function defaultLocalModelInitializer(options: LocalModelInitOptions): Promise { const { FlagEmbedding } = await loadFastembed(); try { return await FlagEmbedding.init(options); } catch (error) { const message = error instanceof Error ? error.message : ""; + if (/Protobuf parsing failed/i.test(message) && (await quarantineCorruptModelFile(message))) { + return FlagEmbedding.init(options); + } if ( !/(?:Config file not found at .*config|Tokenizer file not found at .*tokenizer|Tokens map file not found at .*special_tokens_map)/u.test( message, diff --git a/packages/mnemopi/test/corrupt-model-quarantine.test.ts b/packages/mnemopi/test/corrupt-model-quarantine.test.ts new file mode 100644 index 000000000..760a5abab --- /dev/null +++ b/packages/mnemopi/test/corrupt-model-quarantine.test.ts @@ -0,0 +1,68 @@ +// Contract: a "Protobuf parsing failed" init error quarantines EXACTLY the +// model file named in the message (atomic rename to *.corrupt-) and +// reports retry-safety; unrelated init errors never touch the filesystem. +import { describe, expect, test } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { quarantineCorruptModelFile } from "../src/core/embeddings"; + +async function tempModelFile(): Promise { + const dir = await fs.mkdtemp(path.join(os.tmpdir(), "mnemopi-quarantine-")); + const file = path.join(dir, "model_optimized.onnx"); + await fs.writeFile(file, "not a protobuf"); + return file; +} + +/** The helper only honors paths inside the given cache root. */ +function cacheRootOf(file: string): string { + return path.dirname(path.dirname(file)); +} + +describe("quarantineCorruptModelFile", () => { + test("renames the exact file named by a protobuf failure and allows retry", async () => { + const file = await tempModelFile(); + const healed = await quarantineCorruptModelFile( + `Load model from ${file} failed:Protobuf parsing failed.`, + cacheRootOf(file), + ); + expect(healed).toBe(true); + // Original gone, quarantined copy present. + await expect(fs.access(file)).rejects.toThrow(); + const siblings = await fs.readdir(path.dirname(file)); + expect(siblings.some(name => name.startsWith("model_optimized.onnx.corrupt-"))).toBe(true); + await fs.rm(path.dirname(file), { recursive: true, force: true }); + }); + + test("does not treat unrelated init errors as corruption", async () => { + const file = await tempModelFile(); + const healed = await quarantineCorruptModelFile(`Model file not found at ${file}`, cacheRootOf(file)); + expect(healed).toBe(false); + // Untouched: no rename happened. + expect(await fs.access(file).then(() => true)).toBe(true); + await fs.rm(path.dirname(file), { recursive: true, force: true }); + }); + + test("a missing file (concurrent heal) still reports retry-safe", async () => { + const ghost = path.join(os.tmpdir(), `mnemopi-ghost-${Date.now()}`, "model_optimized.onnx"); + const healed = await quarantineCorruptModelFile( + `Load model from ${ghost} failed:Protobuf parsing failed.`, + path.dirname(path.dirname(ghost)), + ); + expect(healed).toBe(true); + }); + + test("refuses to touch a file OUTSIDE the fastembed cache directory", async () => { + const file = await tempModelFile(); + // Cache root that does NOT contain the file: containment must reject. + const foreignRoot = await fs.mkdtemp(path.join(os.tmpdir(), "mnemopi-foreign-")); + const healed = await quarantineCorruptModelFile( + `Load model from ${file} failed:Protobuf parsing failed.`, + foreignRoot, + ); + expect(healed).toBe(false); + expect(await fs.access(file).then(() => true)).toBe(true); + await fs.rm(path.dirname(file), { recursive: true, force: true }); + await fs.rm(foreignRoot, { recursive: true, force: true }); + }); +});