@danielblomma/cortex-mcp 2.2.1 → 2.2.3

package/scaffold/mcp/src/search.ts

@@ -2,7 +2,7 @@ import { embedQuery, getEmbeddingRuntimeWarning, loadEmbeddingIndex } from "./em
 import {
   buildChunkPartOfRelations,
   buildEntitySearchMap,
-  buildSearchEntities,
+  iterateSearchEntities,
   entityCatalog
 } from "./contextEntities.js";
 import { relationDegree } from "./graphMetrics.js";
@@ -19,7 +19,7 @@ import {
   semanticScore,
   tokenize
 } from "./searchCore.js";
-import { buildSearchResults } from "./searchResults.js";
+import { buildSearchResultsWithStats } from "./searchResults.js";
 import { traverseImpactGraph } from "./impactTraversal.js";
 import { buildEmptyRelatedResponse, buildRelatedResponseMeta } from "./relatedResponse.js";
 import { traverseRelatedGraph } from "./relatedTraversal.js";
@@ -39,6 +39,7 @@ import type {
   ImpactParams,
   RelatedParams,
   RelationType,
+  SearchEntity,
   SearchParams,
   ToolPayload
 } from "./types.js";
@@ -57,6 +58,17 @@ const IMPACT_RELATION_TYPES = new Set([
   "PART_OF"
 ]);
 
+function* filterSearchCandidates(
+  candidates: Iterable<SearchEntity>,
+  includeDeprecated: boolean
+): Generator<SearchEntity> {
+  for (const entity of candidates) {
+    if (includeDeprecated || entity.status.toLowerCase() !== "deprecated") {
+      yield entity;
+    }
+  }
+}
+
 export async function runContextSearch(parsed: SearchParams): Promise<ToolPayload> {
   const searchPresetConfig = resolveSearchResponsePreset(parsed);
   const responsePreset = searchPresetConfig.responsePreset;
@@ -68,16 +80,15 @@ export async function runContextSearch(parsed: SearchParams): Promise<ToolPayloa
   const degreeByEntity = relationDegree(allRelations);
   const queryTokens = expandQueryTokens(Array.from(new Set(tokenize(parsed.query))));
   const queryPhrase = normalizeText(parsed.query).trim();
-  const candidates = buildSearchEntities(data, includeContent).filter(
-    (entity) => parsed.include_deprecated || entity.status.toLowerCase() !== "deprecated"
-  );
+  const candidates = () =>
+    filterSearchCandidates(iterateSearchEntities(data, includeContent), parsed.include_deprecated);
   const embeddings = loadEmbeddingIndex();
   const queryVector =
     embeddings.model && embeddings.vectors.size > 0
       ? await embedQuery(parsed.query, embeddings.model)
       : null;
 
-  const results = buildSearchResults({
+  const { results, totalCandidates } = buildSearchResultsWithStats({
     candidates,
     degreeByEntity,
     queryTokens,
@@ -107,7 +118,7 @@ export async function runContextSearch(parsed: SearchParams): Promise<ToolPayloa
     include_matched_rules: includeMatchedRules,
     include_content: includeContent,
     ranking: data.ranking,
-    total_candidates: candidates.length,
+    total_candidates: totalCandidates,
     context_source: data.source,
     warning: warningMessages.length > 0 ? warningMessages.join(" | ") : undefined,
     semantic_engine:

package/scaffold/mcp/src/searchResults.ts

@@ -1,5 +1,19 @@
 import type { RankingWeights, SearchEntity } from "./types.js";
 
+type SearchResult = Record<string, unknown>;
+
+type RankedSearchResult = {
+  result: SearchResult;
+  rankScore: number;
+};
+
+type SearchCandidateSource = Iterable<SearchEntity> | (() => Iterable<SearchEntity>);
+
+type BaseChunkMetadata = {
+  label: string;
+  path: string;
+};
+
 function isWindowChunkId(id: string): boolean {
   return id.includes(":window:");
 }
@@ -14,8 +28,26 @@ function baseChunkLabel(label: string): string {
   return markerIndex === -1 ? label : label.slice(0, markerIndex);
 }
 
-export function buildSearchResults(params: {
-  candidates: SearchEntity[];
+function pruneRankedResults(bestById: Map<string, RankedSearchResult>, limit: number): void {
+  if (bestById.size <= limit) {
+    return;
+  }
+
+  const retained = [...bestById.entries()]
+    .sort(([, a], [, b]) => b.rankScore - a.rankScore)
+    .slice(0, limit);
+  bestById.clear();
+  for (const [id, result] of retained) {
+    bestById.set(id, result);
+  }
+}
+
+function candidatesFrom(source: SearchCandidateSource): Iterable<SearchEntity> {
+  return typeof source === "function" ? source() : source;
+}
+
+export function buildSearchResultsWithStats(params: {
+  candidates: SearchCandidateSource;
   degreeByEntity: Map<string, number>;
   queryTokens: string[];
   queryPhrase: string;
@@ -32,7 +64,11 @@ export function buildSearchResults(params: {
   vectorScorer: (a: Float32Array, b: Float32Array) => number;
   recencyScorer: (isoDate: string) => number;
   legacyDataAccessBooster: (entity: SearchEntity, queryTokens: string[], queryPhrase: string) => number;
-}): Record<string, unknown>[] {
+}): { results: Record<string, unknown>[]; totalCandidates: number } {
+  if (params.topK <= 0) {
+    return { results: [], totalCandidates: 0 };
+  }
+
   // Graph score = midrank percentile of relation degree within the entity's
   // own type. The previous min(1, degree/4) saturated at degree >= 4, which
   // nearly every entity exceeds, making the graph weight a constant. Per-type
@@ -40,7 +76,10 @@ export function buildSearchResults(params: {
   // (every type averages ~0.5), so hub-heavy types like rules cannot drown
   // out leaf code files or doc sections.
   const sortedDegreesByType = new Map<string, number[]>();
-  for (const entity of params.candidates) {
+  const chunkCandidatesById = new Map<string, BaseChunkMetadata>();
+  let totalCandidates = 0;
+  for (const entity of candidatesFrom(params.candidates)) {
+    totalCandidates += 1;
     const degree = params.degreeByEntity.get(entity.id) ?? 0;
     const list = sortedDegreesByType.get(entity.entity_type);
     if (list) {
@@ -48,6 +87,13 @@ export function buildSearchResults(params: {
     } else {
       sortedDegreesByType.set(entity.entity_type, [degree]);
     }
+
+    if (entity.entity_type === "Chunk" && !isWindowChunkId(entity.id)) {
+      chunkCandidatesById.set(entity.id, {
+        label: entity.label,
+        path: entity.path
+      });
+    }
   }
   for (const list of sortedDegreesByType.values()) {
     list.sort((a, b) => a - b);
@@ -65,102 +111,100 @@ export function buildSearchResults(params: {
     return sorted.length > 0 ? (lo + upper) / (2 * sorted.length) : 0;
   };
 
-  const rawResults = params.candidates
-    .map((entity) => {
-      const lexicalSemantic = params.semanticScorer(params.queryTokens, params.queryPhrase, entity.text);
-      const entityVector = params.embeddingVectors.get(entity.id);
-      const vectorSemantic =
-        params.queryVector && entityVector
-          ? Math.max(0, Math.min(1, params.vectorScorer(params.queryVector, entityVector)))
-          : 0;
-      const hasRelevanceSignal =
-        lexicalSemantic >= params.minLexicalRelevance || vectorSemantic >= params.minVectorRelevance;
-      if (!hasRelevanceSignal) {
-        return null;
-      }
+  const resultLimit = Math.max(1, params.topK);
+  const pruneThreshold = Math.max(resultLimit * 4, 64);
+  const bestById = new Map<string, RankedSearchResult>();
 
-      const semantic =
-        vectorSemantic > 0 ? vectorSemantic * 0.75 + lexicalSemantic * 0.25 : lexicalSemantic;
-      const degree = params.degreeByEntity.get(entity.id) ?? 0;
-      const graphScore = midrankPercentile(sortedDegreesByType.get(entity.entity_type) ?? [], degree);
-      const trustScore = Math.max(0, Math.min(1, entity.trust_level / 100));
-      const dateScore = params.recencyScorer(entity.updated_at);
-
-      let score = 0;
-      score += params.ranking.semantic * semantic;
-      score += params.ranking.graph * graphScore;
-      score += params.ranking.trust * trustScore;
-      score += params.ranking.recency * dateScore;
-      score += params.legacyDataAccessBooster(entity, params.queryTokens, params.queryPhrase);
-
-      if (entity.source_of_truth) {
-        score += 0.1 * semantic;
-      }
+  for (const entity of candidatesFrom(params.candidates)) {
+    const lexicalSemantic = params.semanticScorer(params.queryTokens, params.queryPhrase, entity.text);
+    const entityVector = params.embeddingVectors.get(entity.id);
+    const vectorSemantic =
+      params.queryVector && entityVector
+        ? Math.max(0, Math.min(1, params.vectorScorer(params.queryVector, entityVector)))
+        : 0;
+    const hasRelevanceSignal =
+      lexicalSemantic >= params.minLexicalRelevance || vectorSemantic >= params.minVectorRelevance;
+    if (!hasRelevanceSignal) {
+      continue;
+    }
 
-      return {
-        id: entity.id,
-        entity_type: entity.entity_type,
-        kind: entity.kind,
-        title: entity.label,
-        path: entity.path || undefined,
-        source_of_truth: entity.source_of_truth,
-        status: entity.status,
-        updated_at: entity.updated_at,
-        excerpt: entity.snippet,
-        ...(params.includeScores
-          ? {
-              score: Number(score.toFixed(4)),
-              semantic_score: Number(semantic.toFixed(4)),
-              embedding_score: Number(vectorSemantic.toFixed(4)),
-              lexical_score: Number(lexicalSemantic.toFixed(4)),
-              graph_score: Number(graphScore.toFixed(4))
-            }
-          : {}),
-        ...(params.includeMatchedRules
-          ? {
-              matched_rules: entity.matched_rules
-            }
-          : {}),
-        ...(params.includeContent
-          ? {
-              content: entity.content
-            }
-          : {})
-      };
-    })
-    .filter((result): result is NonNullable<typeof result> => result !== null)
-    .sort((a, b) => Number(b.score ?? 0) - Number(a.score ?? 0));
-
-  const chunkCandidatesById = new Map(
-    params.candidates.filter((entity) => entity.entity_type === "Chunk").map((entity) => [entity.id, entity])
-  );
-  const normalizedById = new Map<string, (typeof rawResults)[number]>();
-
-  for (const result of rawResults) {
-    if (result.entity_type !== "Chunk" || !isWindowChunkId(String(result.id))) {
-      if (!normalizedById.has(String(result.id))) {
-        normalizedById.set(String(result.id), result);
+    const semantic =
+      vectorSemantic > 0 ? vectorSemantic * 0.75 + lexicalSemantic * 0.25 : lexicalSemantic;
+    const degree = params.degreeByEntity.get(entity.id) ?? 0;
+    const graphScore = midrankPercentile(sortedDegreesByType.get(entity.entity_type) ?? [], degree);
+    const trustScore = Math.max(0, Math.min(1, entity.trust_level / 100));
+    const dateScore = params.recencyScorer(entity.updated_at);
+
+    let score = 0;
+    score += params.ranking.semantic * semantic;
+    score += params.ranking.graph * graphScore;
+    score += params.ranking.trust * trustScore;
+    score += params.ranking.recency * dateScore;
+    score += params.legacyDataAccessBooster(entity, params.queryTokens, params.queryPhrase);
+
+    if (entity.source_of_truth) {
+      score += 0.1 * semantic;
+    }
+
+    const result: SearchResult = {
+      id: entity.id,
+      entity_type: entity.entity_type,
+      kind: entity.kind,
+      title: entity.label,
+      path: entity.path || undefined,
+      source_of_truth: entity.source_of_truth,
+      status: entity.status,
+      updated_at: entity.updated_at,
+      excerpt: entity.snippet,
+      ...(params.includeScores
+        ? {
+            score: Number(score.toFixed(4)),
+            semantic_score: Number(semantic.toFixed(4)),
+            embedding_score: Number(vectorSemantic.toFixed(4)),
+            lexical_score: Number(lexicalSemantic.toFixed(4)),
+            graph_score: Number(graphScore.toFixed(4))
+          }
+        : {}),
+      ...(params.includeMatchedRules
+        ? {
+            matched_rules: entity.matched_rules
+          }
+        : {}),
+      ...(params.includeContent
+        ? {
+            content: entity.content
+          }
+        : {})
+    };
+
+    if (entity.entity_type === "Chunk" && isWindowChunkId(entity.id)) {
+      const canonicalId = baseChunkId(entity.id);
+      const baseChunk = chunkCandidatesById.get(canonicalId);
+      if (baseChunk) {
+        result.id = canonicalId;
+        result.title = baseChunkLabel(baseChunk.label);
+        result.path = baseChunk.path || undefined;
       }
-      continue;
     }
 
-    const canonicalId = baseChunkId(String(result.id));
-    const baseChunk = chunkCandidatesById.get(canonicalId);
-    const normalizedResult = baseChunk
-      ? {
-          ...result,
-          id: canonicalId,
-          title: baseChunkLabel(baseChunk.label),
-          path: baseChunk.path || undefined
-        }
-      : result;
-    const existing = normalizedById.get(String(normalizedResult.id));
-    if (!existing || Number(normalizedResult.score ?? 0) > Number(existing.score ?? 0)) {
-      normalizedById.set(String(normalizedResult.id), normalizedResult);
+    const resultId = String(result.id);
+    const existing = bestById.get(resultId);
+    if (!existing || score > existing.rankScore) {
+      bestById.set(resultId, { result, rankScore: score });
+    }
+
+    if (bestById.size > pruneThreshold) {
+      pruneRankedResults(bestById, resultLimit);
     }
   }
 
-  return [...normalizedById.values()]
-    .sort((a, b) => Number(b.score ?? 0) - Number(a.score ?? 0))
-    .slice(0, params.topK);
+  const results = [...bestById.values()]
+    .sort((a, b) => b.rankScore - a.rankScore)
+    .slice(0, params.topK)
+    .map((ranked) => ranked.result);
+  return { results, totalCandidates };
+}
+
+export function buildSearchResults(params: Parameters<typeof buildSearchResultsWithStats>[0]): Record<string, unknown>[] {
+  return buildSearchResultsWithStats(params).results;
 }

package/scaffold/mcp/tests/search-graph-score.test.mjs

@@ -1,6 +1,6 @@
 import test from "node:test";
 import assert from "node:assert/strict";
-import { buildSearchResults } from "../dist/searchResults.js";
+import { buildSearchResults, buildSearchResultsWithStats } from "../dist/searchResults.js";
 
 function makeEntity(id, entityType, overrides = {}) {
   return {
@@ -130,3 +130,76 @@ test("graph score never saturates to a shared constant for common degrees", () =
 
   assert.equal(unique.size, 3);
 });
+
+test("ranking is preserved when score fields are omitted from the response", () => {
+  const candidates = [
+    makeEntity("weak", "Chunk", { text: "weak query" }),
+    makeEntity("strong", "Chunk", { text: "strong query" }),
+    makeEntity("middle", "Chunk", { text: "middle query" })
+  ];
+
+  const results = buildSearchResults({
+    candidates,
+    degreeByEntity: new Map(),
+    queryTokens: ["query"],
+    queryPhrase: "query",
+    ranking: { semantic: 1, graph: 0, trust: 0, recency: 0 },
+    includeScores: false,
+    includeMatchedRules: false,
+    includeContent: false,
+    queryVector: null,
+    embeddingVectors: new Map(),
+    topK: 2,
+    minLexicalRelevance: 0,
+    minVectorRelevance: 0,
+    semanticScorer: (_tokens, _phrase, text) => {
+      if (text.includes("strong")) return 0.9;
+      if (text.includes("middle")) return 0.5;
+      return 0.1;
+    },
+    vectorScorer: () => 0,
+    recencyScorer: () => 0,
+    legacyDataAccessBooster: () => 0
+  });
+
+  assert.deepEqual(results.map((result) => result.id), ["strong", "middle"]);
+  assert.equal("score" in results[0], false);
+});
+
+test("search results can consume a repeatable candidate generator with stats", () => {
+  let passes = 0;
+  const candidateSource = function* () {
+    passes += 1;
+    yield makeEntity("weak", "Chunk", { text: "weak query" });
+    yield makeEntity("strong", "Chunk", { text: "strong query" });
+    yield makeEntity("middle", "Chunk", { text: "middle query" });
+  };
+
+  const { results, totalCandidates } = buildSearchResultsWithStats({
+    candidates: candidateSource,
+    degreeByEntity: new Map(),
+    queryTokens: ["query"],
+    queryPhrase: "query",
+    ranking: { semantic: 1, graph: 0, trust: 0, recency: 0 },
+    includeScores: false,
+    includeMatchedRules: false,
+    includeContent: false,
+    queryVector: null,
+    embeddingVectors: new Map(),
+    topK: 2,
+    minLexicalRelevance: 0,
+    minVectorRelevance: 0,
+    semanticScorer: (_tokens, _phrase, text) => {
+      if (text.includes("strong")) return 0.9;
+      if (text.includes("middle")) return 0.5;
+      return 0.1;
+    },
+    vectorScorer: () => 0,
+    recencyScorer: () => 0,
+    legacyDataAccessBooster: () => 0
+  });
+
+  assert.equal(passes, 2);
+  assert.equal(totalCandidates, 3);
+  assert.deepEqual(results.map((result) => result.id), ["strong", "middle"]);
+});

package/scaffold/scripts/ingest.mjs

@@ -758,6 +758,29 @@ function writeJsonl(filePath, records) {
   }
 }
 
+function stageJsonl(filePath, records) {
+  const stagedPath = `${filePath}.tmp-${process.pid}-${Date.now()}`;
+  writeJsonl(stagedPath, records);
+  return {
+    commit() {
+      fs.renameSync(stagedPath, filePath);
+    },
+    discard() {
+      fs.rmSync(stagedPath, { force: true });
+    }
+  };
+}
+
+function countFileContentRecords(fileRecords) {
+  let count = 0;
+  for (const fileRecord of fileRecords) {
+    if (Object.prototype.hasOwnProperty.call(fileRecord, "content")) {
+      count += 1;
+    }
+  }
+  return count;
+}
+
 function sanitizeTsvCell(value) {
   if (value === null || value === undefined) return "";
   return String(value).replace(/\t/g, " ").replace(/\r?\n/g, " ");
@@ -2288,7 +2311,8 @@ function resolveIngestWorkerCount(taskCount) {
   const raw = process.env.CORTEX_INGEST_WORKERS;
   const configured = raw !== undefined ? Number.parseInt(raw, 10) : Number.NaN;
   const cpuBudget = Math.max(1, (os.availableParallelism?.() ?? os.cpus().length) - 1);
-  const desired = Number.isFinite(configured) && configured >= 0 ? configured : Math.min(cpuBudget, 8);
+  const defaultWorkerLimit = taskCount >= 1000 ? 4 : 8;
+  const desired = Number.isFinite(configured) && configured >= 0 ? configured : Math.min(cpuBudget, defaultWorkerLimit);
   if (desired <= 1) return 1;
   // Worker spin-up plus per-worker WASM grammar init dominates on small or
   // incremental runs; stay sequential until there is enough work to amortize.
@@ -3352,6 +3376,12 @@ async function main() {
 
   const constrainsRelations = [];
   const implementsRelations = [];
+  const stagedDocumentCache = stageJsonl(path.join(CACHE_DIR, "documents.jsonl"), fileRecords);
+  const stagedFileEntityCache = stageJsonl(path.join(CACHE_DIR, "entities.file.jsonl"), fileRecords);
+  memoryTrace.checkpoint("writes:file_cache_staged", {
+    files: fileRecords.length,
+    file_content_records: countFileContentRecords(fileRecords)
+  });
   memoryTrace.checkpoint("tokens:rule_matching_start", {
     files: fileRecords.length,
     rules: ruleRecords.length
@@ -3365,9 +3395,10 @@ async function main() {
   const constrainsByKey = new Map();
   const implementsByKey = new Map();
   let fileTokenSetsProcessed = 0;
+  let fileContentRecordsReleased = 0;
 
   for (const fileRecord of fileRecords) {
-    const lower = fileRecord.content.toLowerCase();
+    const lower = String(fileRecord.content ?? "").toLowerCase();
     const tokens = fileTokenSet(fileRecord);
     fileTokenSetsProcessed += 1;
 
@@ -3406,6 +3437,11 @@ async function main() {
         }
       }
     }
+
+    if (Object.prototype.hasOwnProperty.call(fileRecord, "content")) {
+      delete fileRecord.content;
+      fileContentRecordsReleased += 1;
+    }
   }
 
   for (const { ruleRecord } of ruleMatchers) {
@@ -3430,6 +3466,8 @@ async function main() {
   memoryTrace.checkpoint("tokens:rule_matching_complete", {
     file_token_sets: fileTokenSetsProcessed,
     file_token_sets_retained: 0,
+    file_content_records_released: fileContentRecordsReleased,
+    file_content_records_retained: countFileContentRecords(fileRecords),
     rules: ruleRecords.length,
     relations_constrains: constrainsRelations.length,
     relations_implements: implementsRelations.length,
@@ -3442,8 +3480,8 @@ async function main() {
     rules: ruleRecords.length,
     chunks: chunkRecords.length
   });
-  writeJsonl(path.join(CACHE_DIR, "documents.jsonl"), fileRecords);
-  writeJsonl(path.join(CACHE_DIR, "entities.file.jsonl"), fileRecords);
+  stagedDocumentCache.commit();
+  stagedFileEntityCache.commit();
   writeJsonl(path.join(CACHE_DIR, "entities.adr.jsonl"), adrRecords);
   writeJsonl(path.join(CACHE_DIR, "entities.rule.jsonl"), ruleRecords);
   writeJsonl(path.join(CACHE_DIR, "entities.chunk.jsonl"), chunkRecords);
@@ -3834,5 +3872,6 @@ export {
   generateSectionHandlerRelations,
   getChunkParserForExtension,
   parseFilesInWorkers,
+  resolveIngestWorkerCount,
   resolveRelativeImportTargetId
 };