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:
@@ -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 });
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user