@danielblomma/cortex-mcp 2.2.3 → 2.2.4

package/README.md

@@ -413,12 +413,19 @@ 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 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
+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.
 
 ## Limitations
 

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@danielblomma/cortex-mcp",
   "mcpName": "io.github.DanielBlomma/cortex",
-  "version": "2.2.3",
+  "version": "2.2.4",
   "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

@@ -15,7 +15,9 @@ import {
   resolveMemoryHeadroom,
   resolveModelMaxTokens,
   resolvePoolConfig,
+  resolveTokenBudgetChoice,
   runWorkUnits,
+  truncateTextToTokenBudget,
   type EmbedExtractor,
   type MeasuredText,
   type PendingText
@@ -405,6 +407,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,6 +448,14 @@ 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 uniqueSignatures = new Set<string>();
+  for (const entity of entities) {
+    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);
 
@@ -456,8 +474,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 +489,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
       };
@@ -522,6 +541,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 +589,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 +629,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 +663,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}`
   );
   console.log(
     `[embed] entities=${entities.length} embedded=${embedded} reused=${reused} failed=${failed}`

package/scaffold/mcp/src/embedScheduler.ts

@@ -183,9 +183,46 @@ 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" | "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}`
+  };
 }
 
 /**
@@ -216,6 +253,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,48 @@ 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("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 +504,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 });