fix(mnemopi): self-heal a corrupt cached embedding model on init

A truncated model_optimized.onnx (observed live: 5.7MB file from June,
'Protobuf parsing failed' on every load) blocks local embeddings
forever: the downloader treats the existing file as complete, so every
init re-parses the same broken bytes and local recall/retain loses its
embedder. defaultLocalModelInitializer now quarantines the exact file
named by the loader (atomic rename to *.corrupt-<ts>) and retries init
ONCE so the model re-downloads.

The extracted path is error-message content: it is honored only when
it resolves inside the fastembed cache directory, so a dependency
emitting an unexpected message can never rename an arbitrary file.

Contract tests: protobuf failure quarantines exactly the named cache
file and reports retry-safe; unrelated errors touch nothing; a missing
file (concurrent heal) stays retry-safe; a path outside the cache root
is refused untouched.
This commit is contained in:
DarkPhilosophy
2026-07-18 02:20:44 +03:00
parent 48241afcc4
commit 6cc0c71d1d
2 changed files with 101 additions and 0 deletions
+33
View File
@@ -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<string, Vector>({ max: QUERY_CACHE_MAX });
const providerIds = new WeakMap<object, number>();
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<boolean> {
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<LocalEmbeddingModel> {
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,
@@ -0,0 +1,68 @@
// Contract: a "Protobuf parsing failed" init error quarantines EXACTLY the
// model file named in the message (atomic rename to *.corrupt-<ts>) 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<string> {
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 });
});
});