@danielblomma/cortex-mcp 2.2.3 → 2.2.5

package/README.md

@@ -413,12 +413,20 @@ Required npm configuration:
 ## Embedding performance
 
 Embedding generation tunes itself to the machine: the number of parallel
-workers, memory limits for long files, and skip-work caching are all derived
-from the available CPU cores and RAM at run time (container memory limits
-included). No configuration is needed — on a laptop or a CI runner, cortex
-picks safe, fast settings by itself. The one exception worth knowing:
-when several cortex instances share one machine, set `CORTEX_EMBED_THREADS`
-to give each its fair share of cores.
+workers, memory limits for long files, token budget, and skip-work caching are
+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 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 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.3",
+  "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,11 +11,14 @@ import {
   DEFAULT_SCHEDULER_OPTIONS,
   groupDuplicates,
   packWorkUnits,
+  resolveEffectiveTokenBudget,
   resolveInFlightTokens,
   resolveMemoryHeadroom,
   resolveModelMaxTokens,
   resolvePoolConfig,
+  resolveTokenBudgetChoice,
   runWorkUnits,
+  truncateTextToTokenBudget,
   type EmbedExtractor,
   type MeasuredText,
   type PendingText
@@ -405,6 +408,14 @@ function roundVector(values: number[]): number[] {
   return values.map((value) => Number(value.toFixed(6)));
 }
 
+function resolveSignatureProfile(maxTokenCap: number | null): string {
+  return maxTokenCap ? `embed|max_tokens=${maxTokenCap}` : "";
+}
+
+function embeddingSignature(entitySignature: string, profile: string): string {
+  return profile ? hashText(`${profile}|${entitySignature}`) : entitySignature;
+}
+
 function* presentEmbeddingRecords(
   slots: Array<EmbeddingRecord | null>
 ): Generator<EmbeddingRecord> {
@@ -438,8 +449,12 @@ async function main(): Promise<void> {
   const chunks = parseChunkEntities(readJsonl(path.join(CACHE_DIR, "entities.chunk.jsonl")), filePathById);
 
   const entities: SearchEntity[] = [...documents, ...rules, ...adrs, ...modules, ...projects, ...chunks].sort((a, b) => a.id.localeCompare(b.id));
-
-  const existing = parseExistingEmbeddings(readJsonlRecords(EMBEDDINGS_PATH), modelId);
+  const uniqueSignatures = new Set<string>();
+  for (const entity of entities) {
+    uniqueSignatures.add(entity.signature);
+  }
+  const uniqueTextCount = uniqueSignatures.size;
+  uniqueSignatures.clear();
 
   env.cacheDir = MODEL_CACHE_DIR;
   // Total thread budget for embedding. CORTEX_EMBED_THREADS caps it so
@@ -448,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.
@@ -456,8 +494,9 @@ async function main(): Promise<void> {
   let dimensions = 0;
 
   entities.forEach((entity, index) => {
+    const signature = embeddingSignature(entity.signature, signatureProfile);
     const previous = existing.get(entity.id);
-    if (previous && previous.signature === entity.signature && previous.vector.length > 0) {
+    if (previous && previous.signature === signature && previous.vector.length > 0) {
       reused += 1;
       dimensions = dimensions || previous.vector.length;
       slots[index] = {
@@ -470,7 +509,7 @@ async function main(): Promise<void> {
         source_of_truth: entity.source_of_truth,
         trust_level: entity.trust_level,
         updated_at: entity.updated_at,
-        signature: entity.signature,
+        signature,
         model: modelId,
         dimensions: previous.vector.length
       };
@@ -486,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,
@@ -522,6 +553,7 @@ async function main(): Promise<void> {
   // Fully warm cache: nothing to embed, so skip model loading entirely —
   // this is the common repeat-bootstrap / small-update path.
   let embedded = 0;
+  let modelMaxTokensUsed = 0;
   let result: { vectors: Map<number, number[]>; failures: Array<{ index: number; message: string }> } = {
     vectors: new Map(),
     failures: []
@@ -569,11 +601,22 @@ async function main(): Promise<void> {
 
     // Inference truncates at the model max; token counts must too, or one
     // giant file inflates scheduling cost and gate mass far beyond reality.
-    const modelMaxTokens = resolveModelMaxTokens(first.tokenizer?.model_max_length);
+    const modelMaxTokens = resolveModelMaxTokens(
+      first.tokenizer?.model_max_length,
+      tokenBudget.cap ?? undefined
+    );
+    modelMaxTokensUsed = modelMaxTokens;
     const countTokens = createTokenCounter(first.tokenizer, modelMaxTokens);
+    const countRawTokens = createTokenCounter(first.tokenizer, 131072);
 
-    const measured: MeasuredText[] = unique.map((item) => ({ ...item, tokens: countTokens(item.text) }));
+    const measured: MeasuredText[] = unique.map((item) => {
+      const text = truncateTextToTokenBudget(item.text, countRawTokens, modelMaxTokens);
+      return { ...item, text, tokens: countTokens(text) };
+    });
     const units = packWorkUnits(measured, schedulerOptions);
+    console.log(
+      `[embed] scheduler unique=${unique.length} units=${units.length} max_tokens<=${modelMaxTokens} token_budget=${tokenBudget.mode} reason=${tokenBudget.reason}`
+    );
     // Recompute headroom after the model copies are resident so the gate
     // reflects what is actually left for inference activations.
     const inFlightRaw = Number(process.env.CORTEX_EMBED_INFLIGHT_TOKENS);
@@ -598,7 +641,7 @@ async function main(): Promise<void> {
           source_of_truth: entity.source_of_truth,
           trust_level: entity.trust_level,
           updated_at: entity.updated_at,
-          signature: entity.signature,
+          signature: embeddingSignature(entity.signature, signatureProfile),
           model: modelId,
           dimensions: vector.length,
           vector
@@ -632,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}`
+    `[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

@@ -183,9 +183,114 @@ export function resolveMemoryHeadroom({
 /** Sanitizes a tokenizer-reported maximum sequence length. Sentinel and
  * absurd values (some configs report 1e30) fall back to 8192; anything
  * beyond 128k tokens is treated as absurd. */
-export function resolveModelMaxTokens(raw: unknown): number {
+export function resolveModelMaxTokens(raw: unknown, overrideRaw?: unknown): number {
   const value = Number(raw);
-  return Number.isFinite(value) && value >= 1 && value <= 131072 ? Math.floor(value) : 8192;
+  const modelMax = Number.isFinite(value) && value >= 1 && value <= 131072 ? Math.floor(value) : 8192;
+  const override = Number(overrideRaw);
+  if (Number.isFinite(override) && override >= 16 && override <= 131072) {
+    return Math.min(modelMax, Math.floor(override));
+  }
+  return modelMax;
+}
+
+export type TokenBudgetChoice = {
+  cap: number | null;
+  mode: "auto" | "auto_degraded" | "explicit" | "model";
+  reason: string;
+};
+
+/**
+ * Chooses the requested embedding token budget before model load, so cache
+ * signatures can include the same cap that inference will later enforce.
+ * Auto mode is quality-preserving: it uses the model's own maximum by default
+ * and only truncates when the user explicitly requests a numeric cap.
+ */
+export function resolveTokenBudgetChoice(overrideRaw: unknown, uniqueCount: number): TokenBudgetChoice {
+  const raw = String(overrideRaw ?? "").trim().toLowerCase();
+  const override = Number(raw);
+  if (Number.isFinite(override) && override >= 16 && override <= 131072) {
+    const cap = Math.floor(override);
+    return { cap, mode: "explicit", reason: "env_override" };
+  }
+
+  if (raw === "model" || raw === "none" || raw === "off" || raw === "full") {
+    return { cap: null, mode: "model", reason: "env_full_model" };
+  }
+
+  const count = Number.isFinite(uniqueCount) && uniqueCount > 0 ? Math.floor(uniqueCount) : 0;
+  return {
+    cap: null,
+    mode: "auto",
+    reason: `quality_preserving_model_max:unique=${count}`
+  };
+}
+
+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)}`
+  };
 }
 
 /**
@@ -216,6 +321,29 @@ export function createTokenCounter(
   };
 }
 
+export function truncateTextToTokenBudget(
+  text: string,
+  countTokens: (text: string) => number,
+  maxTokens: number
+): string {
+  if (!Number.isFinite(maxTokens) || maxTokens < 1 || countTokens(text) <= maxTokens) {
+    return text;
+  }
+
+  let low = 0;
+  let high = text.length;
+  while (low < high) {
+    const mid = Math.ceil((low + high) / 2);
+    if (countTokens(text.slice(0, mid)) <= maxTokens) {
+      low = mid;
+    } else {
+      high = mid - 1;
+    }
+  }
+
+  return text.slice(0, low).trimEnd();
+}
+
 /** Estimated working memory one full pool session needs (model copy plus
  * inference activations for a base-size model). Used only as a derating
  * heuristic; deliberately conservative. */

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

@@ -448,6 +448,100 @@ test("resolveModelMaxTokens: sane values pass, sentinels and absurdities fall ba
   assert.equal(resolveModelMaxTokens(0), 8192);
 });
 
+test("resolveModelMaxTokens: explicit override caps but never raises model max", async () => {
+  const { resolveModelMaxTokens } = await import("../dist/embedScheduler.js");
+  assert.equal(resolveModelMaxTokens(8192, 2048), 2048);
+  assert.equal(resolveModelMaxTokens(512, 2048), 512);
+  assert.equal(resolveModelMaxTokens(8192, "1024"), 1024);
+  assert.equal(resolveModelMaxTokens(8192, 4), 8192);
+  assert.equal(resolveModelMaxTokens(8192, "nope"), 8192);
+});
+
+test("resolveTokenBudgetChoice: auto preserves model max regardless of repo size", async () => {
+  const { resolveTokenBudgetChoice } = await import("../dist/embedScheduler.js");
+  assert.deepEqual(resolveTokenBudgetChoice(undefined, 180), {
+    cap: null,
+    mode: "auto",
+    reason: "quality_preserving_model_max:unique=180"
+  });
+  assert.deepEqual(resolveTokenBudgetChoice("auto", 20000), {
+    cap: null,
+    mode: "auto",
+    reason: "quality_preserving_model_max:unique=20000"
+  });
+  assert.deepEqual(resolveTokenBudgetChoice(undefined, 48533), {
+    cap: null,
+    mode: "auto",
+    reason: "quality_preserving_model_max:unique=48533"
+  });
+});
+
+test("resolveTokenBudgetChoice: explicit caps and full-model opt-out win over auto", async () => {
+  const { resolveTokenBudgetChoice } = await import("../dist/embedScheduler.js");
+  assert.deepEqual(resolveTokenBudgetChoice("4096", 48533), {
+    cap: 4096,
+    mode: "explicit",
+    reason: "env_override"
+  });
+  assert.deepEqual(resolveTokenBudgetChoice("model", 48533), {
+    cap: null,
+    mode: "model",
+    reason: "env_full_model"
+  });
+});
+
+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] } });
@@ -462,6 +556,13 @@ test("createTokenCounter: uses the tokenizer, clamps at model max, survives fail
   assert.equal(missing("x".repeat(40)), 10);
 });
 
+test("truncateTextToTokenBudget: trims a prefix to the requested token cap", async () => {
+  const { truncateTextToTokenBudget } = await import("../dist/embedScheduler.js");
+  const countTokens = (text) => text.length;
+  assert.equal(truncateTextToTokenBudget("abcdef", countTokens, 3), "abc");
+  assert.equal(truncateTextToTokenBudget("abc", countTokens, 3), "abc");
+});
+
 test("resolveMemoryHeadroom: constrainedMemory 0 means unconstrained (Node sentinel)", async () => {
   const { resolveMemoryHeadroom } = await import("../dist/embedScheduler.js");
   const unconstrained = resolveMemoryHeadroom({ freeMemory: 50e9, totalMemory: 64e9 });