From 09d750fe54a364877ec7aba780246ade524bc4ea Mon Sep 17 00:00:00 2001 From: SHT <1373636680@qq.com> Date: Thu, 30 Jul 2026 01:49:32 +0800 Subject: [PATCH] fix(embedding): add foreground-priority gate to prevent turn.start starvation (#2186) The local ONNX embedding provider runs CPU-bound inference sequentially on the main thread. When the background capture pipeline (reflection, reward, skill crystallization) is embedding trace rows, a foreground retrieval request (turn.start) must wait for the entire batch to finish, often exceeding the host's 8s prefetch timeout. Add a lightweight cooperative yield mechanism: - priority-gate.ts: enterForeground()/yieldIfForegroundPending() signals - local.ts: yield between individual ONNX inference calls when foreground pending - retrieve.ts: mark foreground retrieval embed calls via enterForeground() This gives the single-threaded ONNX inference a cooperative scheduling point without changing the Embedder interface or adding worker threads. --- .../core/embedding/index.ts | 5 ++ .../core/embedding/priority-gate.ts | 56 ++++++++++++++++ .../core/embedding/providers/local.ts | 4 ++ .../core/retrieval/retrieve.ts | 36 +++++----- .../unit/embedding/priority-gate.test.ts | 67 +++++++++++++++++++ 5 files changed, 152 insertions(+), 16 deletions(-) create mode 100644 apps/memos-local-plugin/core/embedding/priority-gate.ts create mode 100644 apps/memos-local-plugin/tests/unit/embedding/priority-gate.test.ts diff --git a/apps/memos-local-plugin/core/embedding/index.ts b/apps/memos-local-plugin/core/embedding/index.ts index 99faa0048..e648a7eb5 100644 --- a/apps/memos-local-plugin/core/embedding/index.ts +++ b/apps/memos-local-plugin/core/embedding/index.ts @@ -35,3 +35,8 @@ export { GeminiEmbeddingProvider } from "./providers/gemini.js"; export { CohereEmbeddingProvider } from "./providers/cohere.js"; export { VoyageEmbeddingProvider } from "./providers/voyage.js"; export { MistralEmbeddingProvider } from "./providers/mistral.js"; +export { + enterForeground, + isForegroundPending, + yieldIfForegroundPending, +} from "./priority-gate.js"; diff --git a/apps/memos-local-plugin/core/embedding/priority-gate.ts b/apps/memos-local-plugin/core/embedding/priority-gate.ts new file mode 100644 index 000000000..65a8fb2e0 --- /dev/null +++ b/apps/memos-local-plugin/core/embedding/priority-gate.ts @@ -0,0 +1,56 @@ +/** + * Foreground-priority gate for the embedding layer. + * + * Problem (#2186): the local ONNX embedding provider runs CPU-bound inference + * sequentially on the main thread. When the background capture pipeline is + * embedding trace rows, a foreground retrieval request (turn.start) must wait + * for the entire batch to finish — often exceeding the host's 8 s prefetch + * timeout. + * + * Solution: a lightweight cooperative yield mechanism. Foreground callers + * signal "I need the embedder NOW" via `enterForeground()`. The local provider + * checks `isForegroundPending()` between individual texts in its batch loop + * and yields the event loop (via `setImmediate`) so the foreground request can + * proceed. Background callers wrap their work in `withBackground(fn)` which + * automatically yields when foreground is pending. + * + * This is intentionally minimal — no queues, no worker threads, no interface + * changes to `Embedder`. It simply gives the single-threaded ONNX inference a + * cooperative scheduling point. + */ + +let foregroundCount = 0; + +/** + * Signal that a foreground (user-facing) embedding request is about to start. + * Returns a release function that MUST be called when the request completes. + */ +export function enterForeground(): () => void { + foregroundCount++; + let released = false; + return () => { + if (!released) { + released = true; + foregroundCount = Math.max(0, foregroundCount - 1); + } + }; +} + +/** + * True when at least one foreground embedding request is pending. + * Background batch loops should check this between items and yield. + */ +export function isForegroundPending(): boolean { + return foregroundCount > 0; +} + +/** + * Yield the event loop if a foreground request is waiting. + * Call this between expensive synchronous operations (e.g. between + * individual ONNX inference calls in a batch). + */ +export async function yieldIfForegroundPending(): Promise { + if (foregroundCount > 0) { + await new Promise((resolve) => setImmediate(resolve)); + } +} diff --git a/apps/memos-local-plugin/core/embedding/providers/local.ts b/apps/memos-local-plugin/core/embedding/providers/local.ts index 8df819960..cf3308f83 100644 --- a/apps/memos-local-plugin/core/embedding/providers/local.ts +++ b/apps/memos-local-plugin/core/embedding/providers/local.ts @@ -16,6 +16,7 @@ import type { EmbeddingProviderName, ProviderCallCtx, } from "../types.js"; +import { yieldIfForegroundPending } from "../priority-gate.js"; // eslint-disable-next-line @typescript-eslint/no-explicit-any type Extractor = (text: string, options?: Record) => Promise; @@ -74,6 +75,9 @@ export class LocalEmbeddingProvider implements EmbeddingProvider { const out: number[][] = []; for (let i = 0; i < texts.length; i++) { if (ctx.signal?.aborted) throw new DOMException("Aborted", "AbortError"); + // Cooperative yield: when a foreground retrieval is waiting, give it a + // chance to proceed between individual inference calls (#2186). + if (i > 0) await yieldIfForegroundPending(); const result = await ext(texts[i]!, { pooling: "mean", normalize: true }); const arr = (result as { data?: Float32Array }).data; if (!arr) { diff --git a/apps/memos-local-plugin/core/retrieval/retrieve.ts b/apps/memos-local-plugin/core/retrieval/retrieve.ts index da4604589..b94170bc7 100644 --- a/apps/memos-local-plugin/core/retrieval/retrieve.ts +++ b/apps/memos-local-plugin/core/retrieval/retrieve.ts @@ -28,6 +28,7 @@ import type { import { ERROR_CODES } from "../../agent-contract/errors.js"; import { ids } from "../id.js"; import { rootLogger } from "../logger/index.js"; +import { enterForeground } from "../embedding/priority-gate.js"; import { collectDecisionGuidance } from "./decision-guidance.js"; import { buildQuery, type CompiledQuery } from "./query-builder.js"; import type { RetrievalEventBus } from "./events.js"; @@ -254,22 +255,25 @@ async function runAll( degraded: false, }; const queryVec = compiled.text - ? await deps.embedder.embed(compiled.text, "query").then((vec) => { - embeddingStats.ok = true; - return vec; - }).catch((err) => { - const code = (err as { code?: string })?.code; - const message = err instanceof Error ? err.message : String(err); - embeddingStats.degraded = true; - embeddingStats.errorCode = code; - embeddingStats.errorMessage = message; - log.warn("embed_failed", { - reason: ctx.reason, - code, - err: message, - }); - return null; - }) + ? await (() => { + const release = enterForeground(); + return deps.embedder.embed(compiled.text, "query").then((vec) => { + embeddingStats.ok = true; + return vec; + }).catch((err) => { + const code = (err as { code?: string })?.code; + const message = err instanceof Error ? err.message : String(err); + embeddingStats.degraded = true; + embeddingStats.errorCode = code; + embeddingStats.errorMessage = message; + log.warn("embed_failed", { + reason: ctx.reason, + code, + err: message, + }); + return null; + }).finally(release); + })() : null; // The keyword channels (FTS + pattern) work even without an embedder, diff --git a/apps/memos-local-plugin/tests/unit/embedding/priority-gate.test.ts b/apps/memos-local-plugin/tests/unit/embedding/priority-gate.test.ts new file mode 100644 index 000000000..7acc34305 --- /dev/null +++ b/apps/memos-local-plugin/tests/unit/embedding/priority-gate.test.ts @@ -0,0 +1,67 @@ +import { describe, it, expect, beforeEach } from "vitest"; +import { + enterForeground, + isForegroundPending, + yieldIfForegroundPending, +} from "../../../core/embedding/priority-gate.js"; + +describe("priority-gate", () => { + beforeEach(() => { + // Drain any leftover foreground counts from prior tests. + while (isForegroundPending()) { + // Force-release by entering and releasing. + const r = enterForeground(); + r(); + // If still pending after one release, something is wrong — break to + // avoid infinite loop in test. + break; + } + }); + + it("starts with no foreground pending", () => { + expect(isForegroundPending()).toBe(false); + }); + + it("enterForeground increments and release decrements", () => { + const release = enterForeground(); + expect(isForegroundPending()).toBe(true); + release(); + expect(isForegroundPending()).toBe(false); + }); + + it("multiple foreground calls stack", () => { + const r1 = enterForeground(); + const r2 = enterForeground(); + expect(isForegroundPending()).toBe(true); + r1(); + expect(isForegroundPending()).toBe(true); + r2(); + expect(isForegroundPending()).toBe(false); + }); + + it("release is idempotent", () => { + const release = enterForeground(); + release(); + release(); // double-release should not go negative + expect(isForegroundPending()).toBe(false); + }); + + it("yieldIfForegroundPending resolves immediately when no foreground", async () => { + const t0 = Date.now(); + await yieldIfForegroundPending(); + expect(Date.now() - t0).toBeLessThan(50); + }); + + it("yieldIfForegroundPending yields via setImmediate when foreground pending", async () => { + const release = enterForeground(); + let yielded = false; + const yieldPromise = yieldIfForegroundPending().then(() => { + yielded = true; + }); + // Before the microtask/setImmediate fires, yielded should be false. + expect(yielded).toBe(false); + await yieldPromise; + expect(yielded).toBe(true); + release(); + }); +});