@danielblomma/cortex-mcp 2.2.4 → 2.2.5

package/README.md

@@ -418,14 +418,15 @@ all derived from the available CPU cores, RAM, and repository size at run time
 (container memory limits included). No configuration is needed — on a laptop or
 a CI runner, cortex picks safe settings by itself.
 
-The embedding token budget defaults to `auto`, which preserves the embedding
-model's own maximum context. Cortex does not lower the default token budget
-just because a repository is large; lower caps can discard semantically useful
-text and are reserved for explicit benchmark or quality experiments. To force a
+The embedding token budget defaults to `auto`, which is quality-first but
+memory-aware. Cortex starts from the embedding model's own maximum context and
+only lowers the cap when the local memory headroom is unlikely to fit that
+model/context combination. Degraded auto runs are explicit in logs, for example
+`token_budget=auto_degraded reason=memory_headroom cap=2048`. To force a
 specific capped run, set `CORTEX_EMBED_MAX_TOKENS` to a number such as `2048`;
-set it to `model` or `full` to make the full-model baseline explicit. When
-several cortex instances share one machine, set `CORTEX_EMBED_THREADS` to give
-each its fair share of cores.
+set it to `model` or `full` to force the full-model baseline even on tight
+machines. When several cortex instances share one machine, set
+`CORTEX_EMBED_THREADS` to give each its fair share of cores.
 
 ## Limitations
 

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@danielblomma/cortex-mcp",
   "mcpName": "io.github.DanielBlomma/cortex",
-  "version": "2.2.4",
+  "version": "2.2.5",
   "description": "Local, repo-scoped context platform for coding assistants. Semantic search, graph relationships, and architectural rule context.",
   "type": "module",
   "author": "Daniel Blomma",

package/scaffold/mcp/src/embed.ts

@@ -11,6 +11,7 @@ import {
   DEFAULT_SCHEDULER_OPTIONS,
   groupDuplicates,
   packWorkUnits,
+  resolveEffectiveTokenBudget,
   resolveInFlightTokens,
   resolveMemoryHeadroom,
   resolveModelMaxTokens,
@@ -453,11 +454,7 @@ async function main(): Promise<void> {
     uniqueSignatures.add(entity.signature);
   }
   const uniqueTextCount = uniqueSignatures.size;
-  const tokenBudget = resolveTokenBudgetChoice(process.env.CORTEX_EMBED_MAX_TOKENS, uniqueTextCount);
   uniqueSignatures.clear();
-  const signatureProfile = resolveSignatureProfile(tokenBudget.cap);
-
-  const existing = parseExistingEmbeddings(readJsonlRecords(EMBEDDINGS_PATH), modelId);
 
   env.cacheDir = MODEL_CACHE_DIR;
   // Total thread budget for embedding. CORTEX_EMBED_THREADS caps it so
@@ -466,6 +463,29 @@ async function main(): Promise<void> {
   const threadsRaw = Number(process.env.CORTEX_EMBED_THREADS);
   const threadBudget =
     Number.isFinite(threadsRaw) && threadsRaw >= 1 ? Math.floor(threadsRaw) : os.cpus().length;
+  const readHeadroom = () =>
+    resolveMemoryHeadroom({
+      freeMemory: os.freemem(),
+      totalMemory: os.totalmem(),
+      constrainedMemory: process.constrainedMemory?.() ?? null,
+      availableMemory: process.availableMemory?.() ?? null
+    });
+  const memoryHeadroom = readHeadroom();
+  const previewPoolConfig = resolvePoolConfig({
+    threadBudget,
+    uniqueCount: uniqueTextCount,
+    memoryBytes: memoryHeadroom
+  });
+  const requestedTokenBudget = resolveTokenBudgetChoice(process.env.CORTEX_EMBED_MAX_TOKENS, uniqueTextCount);
+  const tokenBudget = resolveEffectiveTokenBudget({
+    choice: requestedTokenBudget,
+    modelMaxTokens: resolveModelMaxTokens(undefined, requestedTokenBudget.cap ?? undefined),
+    memoryBytes: memoryHeadroom,
+    sessions: previewPoolConfig.sessions
+  });
+  const signatureProfile = resolveSignatureProfile(tokenBudget.cap);
+
+  const existing = parseExistingEmbeddings(readJsonlRecords(EMBEDDINGS_PATH), modelId);
 
   let reused = 0;
   // Slot per entity keeps output in entity order; failed slots stay null.
@@ -505,14 +525,6 @@ async function main(): Promise<void> {
   // adapts to the machine so no tuning is expected from users. Container
   // limits (cgroups) and platforms that under-report free memory are both
   // handled inside resolveMemoryHeadroom.
-  const readHeadroom = () =>
-    resolveMemoryHeadroom({
-      freeMemory: os.freemem(),
-      totalMemory: os.totalmem(),
-      constrainedMemory: process.constrainedMemory?.() ?? null,
-      availableMemory: process.availableMemory?.() ?? null
-    });
-  const memoryHeadroom = readHeadroom();
   const poolConfig = resolvePoolConfig({
     threadBudget,
     poolOverride: Number(process.env.CORTEX_EMBED_POOL) || null,
@@ -663,7 +675,7 @@ async function main(): Promise<void> {
   fs.writeFileSync(EMBEDDINGS_MANIFEST_PATH, `${JSON.stringify(manifest, null, 2)}\n`, "utf8");
 
   console.log(
-    `[embed] mode=${mode} model=${modelId} dim=${dimensions} pool=${poolConfig.sessions}x${poolConfig.threadsPerSession} batch<=${schedulerOptions.batchMaxItems} max_tokens<=${modelMaxTokensUsed || tokenBudget.cap || "model"} token_budget=${tokenBudget.mode}`
+    `[embed] mode=${mode} model=${modelId} dim=${dimensions} pool=${poolConfig.sessions}x${poolConfig.threadsPerSession} batch<=${schedulerOptions.batchMaxItems} max_tokens<=${modelMaxTokensUsed || tokenBudget.cap || "model"} token_budget=${tokenBudget.mode} reason=${tokenBudget.reason}`
   );
   console.log(
     `[embed] entities=${entities.length} embedded=${embedded} reused=${reused} failed=${failed}`

package/scaffold/mcp/src/embedScheduler.ts

@@ -195,7 +195,7 @@ export function resolveModelMaxTokens(raw: unknown, overrideRaw?: unknown): numb
 
 export type TokenBudgetChoice = {
   cap: number | null;
-  mode: "auto" | "explicit" | "model";
+  mode: "auto" | "auto_degraded" | "explicit" | "model";
   reason: string;
 };
 
@@ -225,6 +225,74 @@ export function resolveTokenBudgetChoice(overrideRaw: unknown, uniqueCount: numb
   };
 }
 
+const AUTO_TOKEN_BUDGET_CANDIDATES = [8192, 4096, 2048];
+const AUTO_TOKEN_BUDGET_MIN = 2048;
+const AUTO_TOKEN_BUDGET_SESSION_BYTES = 2.4e9;
+const AUTO_TOKEN_BUDGET_ACTIVATION_BYTES_AT_4096 = 3e9;
+const AUTO_TOKEN_BUDGET_MEMORY_MARGIN = 0.95;
+
+function estimateTokenBudgetMemoryBytes(maxTokens: number, sessions: number): number {
+  const sessionCount = Number.isFinite(sessions) && sessions >= 1 ? Math.floor(sessions) : 1;
+  const tokenCount = Number.isFinite(maxTokens) && maxTokens >= 1 ? Math.floor(maxTokens) : 8192;
+  const activationBytes =
+    AUTO_TOKEN_BUDGET_ACTIVATION_BYTES_AT_4096 * Math.pow(tokenCount / 4096, 2);
+  return sessionCount * AUTO_TOKEN_BUDGET_SESSION_BYTES + activationBytes;
+}
+
+/**
+ * Keeps auto quality-first but avoids choosing a token budget that is likely
+ * to OOM this process. Explicit numeric caps and explicit full-model modes are
+ * operator choices and are never degraded here.
+ */
+export function resolveEffectiveTokenBudget({
+  choice,
+  modelMaxTokens,
+  memoryBytes,
+  sessions
+}: {
+  choice: TokenBudgetChoice;
+  modelMaxTokens: number;
+  memoryBytes?: number;
+  sessions?: number;
+}): TokenBudgetChoice {
+  if (choice.mode !== "auto" || choice.cap !== null) {
+    return choice;
+  }
+  const modelMax =
+    Number.isFinite(modelMaxTokens) && modelMaxTokens >= 1 ? Math.floor(modelMaxTokens) : 8192;
+  if (modelMax <= AUTO_TOKEN_BUDGET_MIN) {
+    return choice;
+  }
+  if (!Number.isFinite(memoryBytes) || (memoryBytes as number) <= 0) {
+    return choice;
+  }
+
+  const available = (memoryBytes as number) * AUTO_TOKEN_BUDGET_MEMORY_MARGIN;
+  const candidates = AUTO_TOKEN_BUDGET_CANDIDATES
+    .filter((candidate) => candidate <= modelMax)
+    .sort((a, b) => b - a);
+  if (!candidates.includes(AUTO_TOKEN_BUDGET_MIN)) {
+    candidates.push(AUTO_TOKEN_BUDGET_MIN);
+  }
+
+  const selected =
+    candidates.find((candidate) => estimateTokenBudgetMemoryBytes(candidate, sessions ?? 1) <= available) ??
+    AUTO_TOKEN_BUDGET_MIN;
+
+  if (selected >= modelMax) {
+    return choice;
+  }
+
+  return {
+    cap: selected,
+    mode: "auto_degraded",
+    reason: `memory_headroom cap=${selected} model_max=${modelMax} sessions=${Math.max(
+      1,
+      Math.floor(sessions ?? 1)
+    )} headroom_mb=${Math.round((memoryBytes as number) / 1024 / 1024)}`
+  };
+}
+
 /**
  * Builds a token counter from a tokenizer callable. Falls back to a
  * chars/4 estimate when the tokenizer is unavailable, misbehaves, or

package/scaffold/mcp/tests/embed-scheduler.test.mjs

@@ -490,6 +490,58 @@ test("resolveTokenBudgetChoice: explicit caps and full-model opt-out win over au
   });
 });
 
+test("resolveEffectiveTokenBudget: auto degrades when model max is too expensive for headroom", async () => {
+  const { resolveEffectiveTokenBudget, resolveTokenBudgetChoice } = await import("../dist/embedScheduler.js");
+  const choice = resolveTokenBudgetChoice(undefined, 2237);
+  const degraded = resolveEffectiveTokenBudget({
+    choice,
+    modelMaxTokens: 8192,
+    memoryBytes: 7e9,
+    sessions: 2
+  });
+  assert.equal(degraded.mode, "auto_degraded");
+  assert.equal(degraded.cap, 2048);
+  assert.match(degraded.reason, /memory_headroom/);
+});
+
+test("resolveEffectiveTokenBudget: auto keeps model max when headroom is sufficient", async () => {
+  const { resolveEffectiveTokenBudget, resolveTokenBudgetChoice } = await import("../dist/embedScheduler.js");
+  const choice = resolveTokenBudgetChoice("auto", 2237);
+  assert.deepEqual(
+    resolveEffectiveTokenBudget({
+      choice,
+      modelMaxTokens: 8192,
+      memoryBytes: 64e9,
+      sessions: 2
+    }),
+    choice
+  );
+});
+
+test("resolveEffectiveTokenBudget: explicit caps and full-model opt-out are never degraded", async () => {
+  const { resolveEffectiveTokenBudget, resolveTokenBudgetChoice } = await import("../dist/embedScheduler.js");
+  const explicit = resolveTokenBudgetChoice("4096", 2237);
+  assert.deepEqual(
+    resolveEffectiveTokenBudget({
+      choice: explicit,
+      modelMaxTokens: 8192,
+      memoryBytes: 1e9,
+      sessions: 4
+    }),
+    explicit
+  );
+  const full = resolveTokenBudgetChoice("full", 2237);
+  assert.deepEqual(
+    resolveEffectiveTokenBudget({
+      choice: full,
+      modelMaxTokens: 8192,
+      memoryBytes: 1e9,
+      sessions: 4
+    }),
+    full
+  );
+});
+
 test("createTokenCounter: uses the tokenizer, clamps at model max, survives failures", async () => {
   const { createTokenCounter } = await import("../dist/embedScheduler.js");
   const tokenizer = (text) => ({ input_ids: { dims: [1, text.length * 2] } });