diff --git a/README.md b/README.md index 5b8e90ef..61aed6ce 100644 --- a/README.md +++ b/README.md @@ -280,21 +280,22 @@ the user profile to be injected when enabled. Embeddings power similarity search for memories and the user profile. Configure them in the same file (`~/.config/opencode/opencode-mem.jsonc`). There is **no MLX backend** — local embeddings use `@huggingface/transformers` with ONNX, not Apple MLX. -**Local (default):** set only `embeddingModel`. On first use the model is downloaded from Hugging Face and cached under `{storagePath}/.cache` (default `~/.opencode-mem/data/.cache`). +**Local (default):** set only `embeddingModel`. On first use the model is downloaded from Hugging Face and cached under `{storagePath}/.cache` (default `~/.opencode-mem/data/.cache`). The weights stay resident for the life of the host process; set `embeddingIdleTimeoutMs` to release them after that many idle milliseconds, and the next memory operation re-warms them from the local cache (about a second, not a re-download). **Remote (OpenAI-compatible):** set both `embeddingApiUrl` and `embeddingApiKey`. The plugin then calls `{embeddingApiUrl}/embeddings` with a Bearer token. `embeddingApiKey` accepts the same secret formats as `memoryApiKey` (`literal`, `env://…`, `file://…`). -| Key | Role | -| -------------------------- | ------------------------------------------------------------------------------------------------------------------------------- | -| `embeddingModel` | Hugging Face id (local) or API model name (remote). Default: `Xenova/nomic-embed-text-v1` | -| `embeddingDimensions` | Optional override; usually omit — dimensions are looked up from a built-in map | -| `embeddingPooling` | Local pooling: `"mean"` (default), `"cls"`, or `"last_token"`. Unset → small known-model preset or `"mean"` | -| `embeddingQueryPrefix` | Prefix for query-task embeddings. Unset → model preset or Nomic when `embeddingUseTaskPrefixes` is true. Explicit `""` disables | -| `embeddingDocumentPrefix` | Prefix for document-task embeddings (same resolution rules as `embeddingQueryPrefix`) | -| `embeddingUseTaskPrefixes` | Opt-in Nomic `search_query:` / `search_document:` prefixes when custom/preset prefixes are unset. Default `false` | -| `embeddingDtype` | Optional local ONNX dtype override (e.g. `"q8"`, `"fp32"`). When set, passed to transformers.js `pipeline({ dtype })` | -| `embeddingApiUrl` | Base URL for an OpenAI-compatible embeddings API (no trailing path beyond `/v1`) | -| `embeddingApiKey` | API key for that endpoint (required together with `embeddingApiUrl`) | +| Key | Role | +| -------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------ | +| `embeddingModel` | Hugging Face id (local) or API model name (remote). Default: `Xenova/nomic-embed-text-v1` | +| `embeddingDimensions` | Optional override; usually omit — dimensions are looked up from a built-in map | +| `embeddingPooling` | Local pooling: `"mean"` (default), `"cls"`, or `"last_token"`. Unset → small known-model preset or `"mean"` | +| `embeddingQueryPrefix` | Prefix for query-task embeddings. Unset → model preset or Nomic when `embeddingUseTaskPrefixes` is true. Explicit `""` disables | +| `embeddingDocumentPrefix` | Prefix for document-task embeddings (same resolution rules as `embeddingQueryPrefix`) | +| `embeddingUseTaskPrefixes` | Opt-in Nomic `search_query:` / `search_document:` prefixes when custom/preset prefixes are unset. Default `false` | +| `embeddingDtype` | Optional local ONNX dtype override (e.g. `"q8"`, `"fp32"`). When set, passed to transformers.js `pipeline({ dtype })` | +| `embeddingIdleTimeoutMs` | Release the local model after this many idle milliseconds (e.g. `600000` for 10 minutes). `0` (default) keeps it loaded for the process lifetime | +| `embeddingApiUrl` | Base URL for an OpenAI-compatible embeddings API (no trailing path beyond `/v1`) | +| `embeddingApiKey` | API key for that endpoint (required together with `embeddingApiUrl`) | Recommended local models: diff --git a/src/config.ts b/src/config.ts index 0682a91c..d518f0a0 100644 --- a/src/config.ts +++ b/src/config.ts @@ -65,6 +65,11 @@ interface OpenCodeMemConfig { * When set, passed to `pipeline({ dtype })` and outranks auto-detection. */ embeddingDtype?: string; + /** + * Release the local ONNX pipeline after this many idle milliseconds. + * `0` (default) keeps it loaded for the process lifetime. + */ + embeddingIdleTimeoutMs?: number; similarityThreshold?: number; maxMemories?: number; maxProfileItems?: number; @@ -286,6 +291,7 @@ const DEFAULTS: Required< embeddingDimensions: 768, embeddingUseTaskPrefixes: false, embeddingPooling: "mean", + embeddingIdleTimeoutMs: 0, similarityThreshold: 0.6, maxMemories: 10, maxProfileItems: 5, @@ -444,6 +450,10 @@ const CONFIG_TEMPLATE = `{ // Optional local ONNX dtype override (e.g. "q8", "fp32"). When unset, transformers.js defaults apply. // "embeddingDtype": "q8", + + // Release the local model after this many ms without embedding work. + // 0 (default) keeps it loaded for the whole process lifetime. + // "embeddingIdleTimeoutMs": 600000, // Auto-detected dimensions (no need to set manually) // "embeddingDimensions": 768, @@ -941,6 +951,13 @@ export function normalizeEmbeddingDtype(value: unknown): LocalOnnxDtype | undefi ); } +export function normalizeEmbeddingIdleTimeoutMs(value: number): number { + if (!Number.isInteger(value) || value < 0) { + throw new Error(`Invalid embeddingIdleTimeoutMs config: ${String(value)}`); + } + return value; +} + function getEmbeddingDimensions(model: string): number { const dimensionMap: Record = { // Local Xenova models @@ -1122,6 +1139,9 @@ function buildConfig(fileConfig: OpenCodeMemConfig) { ? fileConfig.embeddingDocumentPrefix : embeddingPreset?.documentPrefix; const embeddingDtype = normalizeEmbeddingDtype(fileConfig.embeddingDtype); + const embeddingIdleTimeoutMs = normalizeEmbeddingIdleTimeoutMs( + fileConfig.embeddingIdleTimeoutMs ?? DEFAULTS.embeddingIdleTimeoutMs + ); const autoCaptureMaxContextBytes = normalizeAutoCaptureMaxContextBytes( fileConfig.autoCaptureMaxContextBytes ?? DEFAULTS.autoCaptureMaxContextBytes ); @@ -1163,6 +1183,7 @@ function buildConfig(fileConfig: OpenCodeMemConfig) { embeddingQueryPrefix, embeddingDocumentPrefix, embeddingDtype, + embeddingIdleTimeoutMs, embeddingApiUrl: fileConfig.embeddingApiUrl, embeddingApiKey: fileConfig.embeddingApiUrl ? resolveSecretValue(fileConfig.embeddingApiKey ?? process.env.OPENAI_API_KEY) diff --git a/src/services/client.ts b/src/services/client.ts index f2448638..feb67807 100644 --- a/src/services/client.ts +++ b/src/services/client.ts @@ -117,6 +117,8 @@ export class LocalMemoryClient { const { releasePluginTurso } = await import("./turso/lifecycle.js"); const closed = await releasePluginTurso(); if (closed) { + // Process-wide like the pool: only the last instance may drop the model. + await embeddingService.dispose("client-close"); this.reset(); } } diff --git a/src/services/embedding.ts b/src/services/embedding.ts index e9065a46..ced6b17b 100644 --- a/src/services/embedding.ts +++ b/src/services/embedding.ts @@ -1,5 +1,5 @@ import { CONFIG } from "../config.js"; -import { log } from "./logger.js"; +import { log, safeLog } from "./logger.js"; import { existsSync, readFileSync } from "node:fs"; import { join } from "node:path"; import { @@ -354,6 +354,12 @@ export class EmbeddingService { public initError: string | null = null; private cache: Map = new Map(); private cachedFingerprint: string | null = null; + private idleTimer: ReturnType | null = null; + private inFlight = 0; + /** Release requested while an embed held the pipeline; runs when it drains. */ + private pendingDispose: string | null = null; + /** Bumped by dispose so a slow initialize cannot resurrect a released pipeline. */ + private generation = 0; static getInstance(): EmbeddingService { if (!(globalThis as any)[GLOBAL_EMBEDDING_KEY]) { @@ -362,6 +368,51 @@ export class EmbeddingService { return (globalThis as any)[GLOBAL_EMBEDDING_KEY]; } + private clearIdleTimer(): void { + if (this.idleTimer !== null) { + clearTimeout(this.idleTimer); + this.idleTimer = null; + } + } + + private scheduleIdleDispose(): void { + this.clearIdleTimer(); + // Remote embeddings hold no model in-process, so there is nothing to release. + if ((CONFIG.embeddingApiUrl && CONFIG.embeddingApiKey) || !this.isWarmedUp) return; + const ms = CONFIG.embeddingIdleTimeoutMs; + if (!(ms > 0)) return; + this.idleTimer = setTimeout(() => { + this.idleTimer = null; + void this.dispose("idle-timeout"); + }, ms); + (this.idleTimer as any).unref?.(); + } + + /** Drop the local pipeline so its weights can be reclaimed; the next embed() re-warms it. */ + async dispose(reason: string = "dispose"): Promise { + this.clearIdleTimer(); + if (this.inFlight > 0) { + // A running embed still holds the pipeline; release once it drains. + this.pendingDispose = reason; + return; + } + this.generation++; + const pipe = this.pipe; + this.pipe = null; + this.isWarmedUp = false; + this.initPromise = null; + this.initError = null; + if (pipe?.dispose) { + try { + await pipe.dispose(); + } catch (error) { + safeLog("Failed to dispose embedding pipeline", { error: String(error) }); + } + } + // Teardown path: a logger failure must not escape the idle timer or close(). + safeLog("Embedding model released", { reason, model: CONFIG.embeddingModel }); + } + async warmup(progressCallback?: (progress: any) => void): Promise { if (this.isWarmedUp) return; if (this.initError) throw new Error(this.initError); @@ -384,6 +435,7 @@ export class EmbeddingService { } // Local model path + const generation = this.generation; const { pipeline } = await ensureTransformersLoaded(); // Explicit config wins; otherwise auto-resolve from model declaration / cache (#393). const dtype = CONFIG.embeddingDtype ?? resolveLocalOnnxDtype(CONFIG.embeddingModel); @@ -394,12 +446,19 @@ export class EmbeddingService { source: CONFIG.embeddingDtype ? "config" : "auto", }); } - this.pipe = await pipeline("feature-extraction", CONFIG.embeddingModel, { + const pipe = await pipeline("feature-extraction", CONFIG.embeddingModel, { ...(dtype ? { dtype } : {}), progress_callback: progressCallback, }); + if (generation !== this.generation) { + // dispose() ran while the model was loading; keep nothing resident. + await (pipe as any).dispose?.(); + return; + } + this.pipe = pipe; this.isWarmedUp = true; this.initError = null; + this.scheduleIdleDispose(); log("Embedding model warmed up", { model: CONFIG.embeddingModel, pooling: CONFIG.embeddingPooling, @@ -433,26 +492,42 @@ export class EmbeddingService { await this.initPromise; } - let result: Float32Array; + this.inFlight++; + try { + let result: Float32Array; - if (CONFIG.embeddingApiUrl && CONFIG.embeddingApiKey) { - const data: any = await fetchRemoteEmbeddingWithRetry(fetch, input); - result = new Float32Array(data.data[0].embedding); - } else { - const output = await this.pipe(input, { - pooling: CONFIG.embeddingPooling, - normalize: true, - }); - result = new Float32Array(output.data); - } + if (CONFIG.embeddingApiUrl && CONFIG.embeddingApiKey) { + const data: any = await fetchRemoteEmbeddingWithRetry(fetch, input); + result = new Float32Array(data.data[0].embedding); + } else { + const pipe = this.pipe; + if (!pipe) { + throw new Error("Embedding model was released before the request ran"); + } + const output = await pipe(input, { + pooling: CONFIG.embeddingPooling, + normalize: true, + }); + result = new Float32Array(output.data); + } - if (this.cache.size >= MAX_CACHE_SIZE) { - const firstKey = this.cache.keys().next().value; - if (firstKey !== undefined) this.cache.delete(firstKey); + if (this.cache.size >= MAX_CACHE_SIZE) { + const firstKey = this.cache.keys().next().value; + if (firstKey !== undefined) this.cache.delete(firstKey); + } + this.cache.set(input, result); + + return result; + } finally { + this.inFlight--; + if (this.pendingDispose) { + const reason = this.pendingDispose; + this.pendingDispose = null; + void this.dispose(reason); + } else { + this.scheduleIdleDispose(); + } } - this.cache.set(input, result); - - return result; } async embedWithTimeout(text: string, options?: EmbedOptions): Promise { diff --git a/tests/config.test.ts b/tests/config.test.ts index c3f794d2..98586ebc 100644 --- a/tests/config.test.ts +++ b/tests/config.test.ts @@ -17,6 +17,7 @@ const { isPlaceholderApiKey, normalizeAutoCleanupRetentionDays, normalizeAutoCaptureMaxContextBytes, + normalizeEmbeddingIdleTimeoutMs, } = await import("../src/config.js"); afterAll(() => { @@ -51,6 +52,19 @@ describe("config", () => { expect(CONFIG.embeddingDtype).toBeUndefined(); }); + it("should keep the local embedding model loaded by default", async () => { + // Cache-busted: CONFIG otherwise carries the developer's real user config. + const fresh = await import(`../src/config.js?cachebust=${Math.random()}`); + expect(fresh.CONFIG.embeddingIdleTimeoutMs).toBe(0); + }); + + it("should normalize embeddingIdleTimeoutMs and allow 0 to disable", () => { + expect(normalizeEmbeddingIdleTimeoutMs(0)).toBe(0); + expect(normalizeEmbeddingIdleTimeoutMs(600000)).toBe(600000); + expect(() => normalizeEmbeddingIdleTimeoutMs(-1)).toThrow(/embeddingIdleTimeoutMs/); + expect(() => normalizeEmbeddingIdleTimeoutMs(1.5)).toThrow(/embeddingIdleTimeoutMs/); + }); + it("should have numeric embeddingDimensions", () => { expect(typeof CONFIG.embeddingDimensions).toBe("number"); expect(CONFIG.embeddingDimensions).toBeGreaterThan(0); diff --git a/tests/embedding-idle-dispose.test.ts b/tests/embedding-idle-dispose.test.ts new file mode 100644 index 00000000..e869b80a --- /dev/null +++ b/tests/embedding-idle-dispose.test.ts @@ -0,0 +1,103 @@ +import { afterAll, describe, expect, it } from "bun:test"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +const home = mkdtempSync(join(tmpdir(), "opencode-mem-idle-")); +const originalHome = process.env.HOME; +const originalUserProfile = process.env.USERPROFILE; +process.env.HOME = home; +process.env.USERPROFILE = home; + +const { CONFIG } = await import("../src/config.js"); +const { EmbeddingService } = await import("../src/services/embedding.js"); + +afterAll(() => { + process.env.HOME = originalHome; + process.env.USERPROFILE = originalUserProfile; +}); + +function makeService(): { service: any; disposed: string[] } { + const disposed: string[] = []; + const service = new EmbeddingService(); + service.pipe = Object.assign(async () => ({ data: new Float32Array([1, 2, 3]) }), { + dispose: async () => { + disposed.push("dispose"); + }, + }); + service.isWarmedUp = true; + return { service, disposed }; +} + +async function withIdleTimeout(ms: number, run: () => Promise): Promise { + const previous = CONFIG.embeddingIdleTimeoutMs; + CONFIG.embeddingIdleTimeoutMs = ms; + CONFIG.embeddingApiUrl = undefined; + CONFIG.embeddingApiKey = undefined; + try { + await run(); + } finally { + CONFIG.embeddingIdleTimeoutMs = previous; + } +} + +describe("EmbeddingService idle release", () => { + it("disposes the local pipeline after the idle timeout", async () => { + await withIdleTimeout(20, async () => { + const { service, disposed } = makeService(); + await service.embed("hello"); + expect(disposed).toEqual([]); + + await new Promise((resolve) => setTimeout(resolve, 80)); + + expect(disposed).toEqual(["dispose"]); + expect(service.isWarmedUp).toBe(false); + expect(service.pipe).toBeNull(); + }); + }); + + it("keeps the model loaded when the timeout is 0", async () => { + await withIdleTimeout(0, async () => { + const { service, disposed } = makeService(); + await service.embed("hello"); + + await new Promise((resolve) => setTimeout(resolve, 80)); + + expect(disposed).toEqual([]); + expect(service.isWarmedUp).toBe(true); + }); + }); + + it("defers disposal while an embed is in flight and releases it afterwards", async () => { + await withIdleTimeout(0, async () => { + const { service, disposed } = makeService(); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + service.pipe = Object.assign( + async () => { + await gate; + return { data: new Float32Array([1]) }; + }, + { dispose: async () => disposed.push("dispose") } + ); + + const running = service.embed("in-flight"); + await new Promise((resolve) => setTimeout(resolve, 10)); + await service.dispose("client-close"); + + // The running embed still holds the pipeline, so nothing is released yet. + expect(disposed).toEqual([]); + + release(); + await running; + await new Promise((resolve) => setTimeout(resolve, 10)); + + // Release happens without a second dispose() call, even with the timer off. + expect(disposed).toEqual(["dispose"]); + expect(service.isWarmedUp).toBe(false); + expect(service.pipe).toBeNull(); + }); + }); +}); diff --git a/tests/fixtures/memory-injection-eligibility.mjs b/tests/fixtures/memory-injection-eligibility.mjs index 60dadcbd..35e18eeb 100644 --- a/tests/fixtures/memory-injection-eligibility.mjs +++ b/tests/fixtures/memory-injection-eligibility.mjs @@ -37,7 +37,7 @@ mock.module(moduleUrl("config"), () => ({ })); mock.module(moduleUrl("services/embedding"), () => ({ EmbeddingService: class {}, - embeddingService: { isWarmedUp: true, warmup: async () => {} }, + embeddingService: { isWarmedUp: true, warmup: async () => {}, dispose: async () => {} }, })); mock.module(moduleUrl("services/turso/ready"), () => ({ ensureTursoReady: async () => {}, diff --git a/tests/plugin-turso-refcount.test.ts b/tests/plugin-turso-refcount.test.ts index e5d36072..529220ab 100644 --- a/tests/plugin-turso-refcount.test.ts +++ b/tests/plugin-turso-refcount.test.ts @@ -78,11 +78,22 @@ describe("plugin Turso refcount (#443)", () => { const id = await userPromptManager.savePrompt("sess-b", "msg-b", "/tmp/p", "world"); + // Fake a warmed local model so the close path has something to release. + const { embeddingService } = await import("../src/services/embedding.js"); + const released: string[] = []; + Object.assign(embeddingService, { + pipe: { dispose: async () => released.push("dispose") }, + isWarmedUp: true, + }); + await memoryClient.close(); expect(pluginTursoRefCount()).toBe(1); expect((await userPromptManager.getPromptById(id))?.content).toBe("world"); + // Sibling instance still live: the process-wide model must stay loaded. + expect(released).toEqual([]); await memoryClient.close(); expect(pluginTursoRefCount()).toBe(0); + expect(released).toEqual(["dispose"]); }); });