forge-server 0.1.0

package/src/index.ts

@@ -0,0 +1,47 @@
+#!/usr/bin/env node
+// index.ts — B19
+// Entry point for the dk-forge-server MCP server.
+// Starts the server on stdio transport for use with Claude and other MCP clients.
+
+import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import { createForgeServer } from './server.js';
+
+async function main(): Promise<void> {
+  const server = await createForgeServer();
+  const transport = new StdioServerTransport();
+
+  // Connect the server to the stdio transport
+  await server.connect(transport);
+
+  // Handle graceful shutdown on SIGINT (Ctrl-C) and SIGTERM
+  const shutdown = async (signal: string) => {
+    process.stderr.write(`\ndk-forge-server: received ${signal}, shutting down...\n`);
+
+    try {
+      // Close storage connections if the server exposes a close method
+      const serverWithClose = server as typeof server & {
+        _closeConnections?: () => Promise<void>;
+      };
+      if (typeof serverWithClose._closeConnections === 'function') {
+        await serverWithClose._closeConnections();
+      }
+
+      await server.close();
+    } catch (err) {
+      process.stderr.write(`dk-forge-server: shutdown error: ${String(err)}\n`);
+    }
+
+    process.exit(0);
+  };
+
+  process.on('SIGINT', () => { void shutdown('SIGINT'); });
+  process.on('SIGTERM', () => { void shutdown('SIGTERM'); });
+
+  // Keep the process alive — stdio transport reads from stdin
+  // Node.js will keep running while stdin is open
+}
+
+main().catch((err) => {
+  process.stderr.write(`dk-forge-server: fatal error: ${String(err)}\n`);
+  process.exit(1);
+});

package/src/ingestion/chunker.ts

@@ -0,0 +1,277 @@
+/**
+ * AST-aware code chunker (ADR-5).
+ * Primary strategy: each top-level entity becomes one chunk (function, class, interface, etc.)
+ * Fallback: fixed-size 512-token chunks with 64-token overlap for non-parseable files.
+ * Large entities (>1024 tokens) are split at method boundaries.
+ */
+
+import { createHash } from 'crypto';
+import { randomUUID } from 'crypto';
+import { extname } from 'path';
+import type { ParsedEntity, ParsedImport, CodeChunk, ChunkType } from '../util/types.js';
+import { estimateTokens } from '../util/token-counter.js';
+
+const CHUNK_MAX_TOKENS = 1024;
+const FIXED_CHUNK_SIZE_TOKENS = 512;
+const FIXED_CHUNK_OVERLAP_TOKENS = 64;
+
+/**
+ * Generate chunks from parsed AST entities.
+ * Called when we have a successful AST parse result.
+ */
+export function chunkFromEntities(
+  entities: ParsedEntity[],
+  imports: ParsedImport[],
+  fileContent: string,
+  filePath: string,
+  repoId: string,
+  language: string,
+  indexedAt: number
+): CodeChunk[] {
+  const chunks: CodeChunk[] = [];
+
+  // Generate the file-level chunk (imports + any top-level statements not in entities)
+  const fileLevelChunk = buildFileLevelChunk(imports, fileContent, filePath, repoId, language, indexedAt, entities);
+  if (fileLevelChunk) chunks.push(fileLevelChunk);
+
+  // Generate one chunk per entity
+  for (const entity of entities) {
+    const entityChunks = buildEntityChunks(entity, filePath, repoId, language, indexedAt);
+    chunks.push(...entityChunks);
+  }
+
+  return chunks;
+}
+
+/**
+ * Generate fixed-size chunks for non-parseable files (markdown, JSON, YAML, config).
+ */
+export function chunkFixed(
+  content: string,
+  filePath: string,
+  repoId: string,
+  language: string,
+  indexedAt: number
+): CodeChunk[] {
+  const ext = extname(filePath).toLowerCase();
+  const chunkType = detectNonCodeType(ext);
+
+  // Split content into lines for better boundary detection
+  const lines = content.split('\n');
+  const chunks: CodeChunk[] = [];
+
+  let currentLines: string[] = [];
+  let currentTokens = 0;
+  let startLine = 1;
+
+  const flushChunk = (endLine: number) => {
+    if (currentLines.length === 0) return;
+    const chunkContent = buildChunkHeader(filePath, null, chunkType, startLine, endLine) +
+      currentLines.join('\n');
+    chunks.push(makeChunk(chunkContent, filePath, repoId, chunkType, startLine, endLine, language, indexedAt));
+  };
+
+  for (const line of lines) {
+    const lineTokens = estimateTokens(line);
+
+    if (currentTokens + lineTokens > FIXED_CHUNK_SIZE_TOKENS && currentLines.length > 0) {
+      flushChunk(startLine + currentLines.length - 1);
+
+      // Overlap: keep last N token-equivalent lines
+      const overlapLines = getOverlapLines(currentLines, FIXED_CHUNK_OVERLAP_TOKENS);
+      startLine = startLine + currentLines.length - overlapLines.length;
+      currentLines = overlapLines;
+      currentTokens = estimateTokensForLines(overlapLines);
+    }
+
+    currentLines.push(line);
+    currentTokens += lineTokens;
+  }
+
+  if (currentLines.length > 0) {
+    flushChunk(startLine + currentLines.length - 1);
+  }
+
+  return chunks;
+}
+
+function buildFileLevelChunk(
+  imports: ParsedImport[],
+  fileContent: string,
+  filePath: string,
+  repoId: string,
+  language: string,
+  indexedAt: number,
+  entities: ParsedEntity[]
+): CodeChunk | null {
+  if (imports.length === 0 && entities.length > 0) return null;
+
+  // Build import block
+  const importLines = imports.map(imp => {
+    if (imp.isDefault) {
+      return `import ${imp.specifiers} from '${imp.fromPath}';`;
+    }
+    return `import ${imp.specifiers} from '${imp.fromPath}';`;
+  });
+
+  // Find the last line of the last import to get file-level lines
+  const importBlock = importLines.join('\n');
+  if (!importBlock.trim()) return null;
+
+  const header = buildChunkHeader(filePath, null, 'file_level', 1, importLines.length);
+  const content = header + importBlock;
+
+  if (estimateTokens(content) < 10) return null;
+
+  return makeChunk(content, filePath, repoId, 'file_level', 1, importLines.length, language, indexedAt);
+}
+
+function buildEntityChunks(
+  entity: ParsedEntity,
+  filePath: string,
+  repoId: string,
+  language: string,
+  indexedAt: number
+): CodeChunk[] {
+  const chunkType = entityTypeToChunkType(entity.type);
+  const header = buildChunkHeader(filePath, entity.name, chunkType, entity.startLine, entity.endLine);
+  const content = header + entity.sourceText;
+
+  const tokens = estimateTokens(content);
+
+  if (tokens <= CHUNK_MAX_TOKENS) {
+    // Single chunk for this entity
+    return [makeChunk(content, filePath, repoId, chunkType, entity.startLine, entity.endLine, language, indexedAt)];
+  }
+
+  // Entity is too large - split at method/property boundaries (heuristic line split)
+  return splitLargeEntity(content, entity, filePath, repoId, chunkType, language, indexedAt);
+}
+
+function splitLargeEntity(
+  content: string,
+  entity: ParsedEntity,
+  filePath: string,
+  repoId: string,
+  chunkType: ChunkType,
+  language: string,
+  indexedAt: number
+): CodeChunk[] {
+  const chunks: CodeChunk[] = [];
+  const lines = content.split('\n');
+  let currentLines: string[] = [];
+  let currentTokens = 0;
+  let chunkIndex = 0;
+  let chunkStartLine = entity.startLine;
+
+  const flush = () => {
+    if (currentLines.length === 0) return;
+    const chunkContent = currentLines.join('\n');
+    const chunkEndLine = chunkStartLine + currentLines.length - 1;
+    chunks.push(makeChunk(
+      chunkContent,
+      filePath, repoId, chunkType,
+      chunkStartLine, chunkEndLine,
+      language, indexedAt,
+      `${entity.name}_part${chunkIndex++}`
+    ));
+    chunkStartLine = chunkEndLine + 1;
+    currentLines = [];
+    currentTokens = 0;
+  };
+
+  for (const line of lines) {
+    const lineTokens = estimateTokens(line);
+
+    if (currentTokens + lineTokens > CHUNK_MAX_TOKENS && currentLines.length > 0) {
+      flush();
+    }
+
+    currentLines.push(line);
+    currentTokens += lineTokens;
+  }
+
+  flush();
+  return chunks;
+}
+
+export function makeChunk(
+  content: string,
+  filePath: string,
+  repoId: string,
+  entityType: ChunkType,
+  startLine: number,
+  endLine: number,
+  language: string,
+  indexedAt: number,
+  entityName?: string
+): CodeChunk {
+  const id = randomUUID();
+  const contentHash = createHash('sha256').update(content, 'utf8').digest('hex');
+
+  return {
+    id,
+    repoId,
+    filePath,
+    entityName: entityName ?? null,
+    entityType,
+    startLine,
+    endLine,
+    language,
+    contentHash,
+    content,
+    contentPreview: content.slice(0, 200),
+    indexedAt,
+  };
+}
+
+export function buildChunkHeader(
+  filePath: string,
+  entityName: string | null,
+  entityType: ChunkType,
+  startLine: number,
+  endLine: number
+): string {
+  const nameStr = entityName ? ` ${entityName}` : '';
+  return `// [dk-forge] file: ${filePath} | entity:${nameStr} (${entityType}) | lines: ${startLine}-${endLine}\n`;
+}
+
+function entityTypeToChunkType(type: string): ChunkType {
+  const map: Record<string, ChunkType> = {
+    function: 'function',
+    class: 'class',
+    interface: 'interface',
+    type_alias: 'type_alias',
+    variable: 'variable',
+    module: 'file_level',
+    file: 'file_level',
+  };
+  return map[type] ?? 'file_level';
+}
+
+function detectNonCodeType(ext: string): ChunkType {
+  if (ext === '.md' || ext === '.mdx') return 'markdown';
+  if (ext === '.json' || ext === '.yaml' || ext === '.yml' || ext === '.toml') return 'config';
+  if (ext === '.graphql' || ext === '.gql') return 'config';
+  if (ext === '.sql') return 'config';
+  return 'config';
+}
+
+function getOverlapLines(lines: string[], targetTokens: number): string[] {
+  let tokens = 0;
+  const result: string[] = [];
+
+  for (let i = lines.length - 1; i >= 0; i--) {
+    const line = lines[i] ?? '';
+    const t = estimateTokens(line);
+    if (tokens + t > targetTokens) break;
+    result.unshift(line);
+    tokens += t;
+  }
+
+  return result;
+}
+
+function estimateTokensForLines(lines: string[]): number {
+  return lines.reduce((sum, line) => sum + estimateTokens(line), 0);
+}

package/src/ingestion/embedder.ts

@@ -0,0 +1,167 @@
+/**
+ * Embedding pipeline using Transformers.js with ONNX backend (ADR-2).
+ * Model: all-MiniLM-L6-v2 (384-dim, ~23 MB quantized, ~5ms/embed).
+ * Singleton pattern - expensive to initialize, reuse across all calls.
+ * Batch embedding: adaptive batch size (32 for long, 64 for short texts).
+ *
+ * Phase 5 changes:
+ *   - Lazy initialization: model loads on first real query, not at startup.
+ *   - isReady() lets callers surface "model loading..." status.
+ *   - Adaptive batching: short texts (<100 chars) use batch size 64.
+ *   - Throughput logging at info level during indexing.
+ */
+
+import { logger } from '../util/logger.js';
+
+const BASE_BATCH_SIZE = 32;
+const SHORT_TEXT_BATCH_SIZE = 64;
+const SHORT_TEXT_THRESHOLD = 100; // chars
+const MODEL_NAME = 'Xenova/all-MiniLM-L6-v2';
+const EMBEDDING_DIMENSION = 384;
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+let pipelineInstance: any = null;
+let initPromise: Promise<void> | null = null;
+let modelReady = false;
+
+/**
+ * Returns true when the embedding pipeline has been initialized and is ready.
+ * False means the first embed call will trigger model load (adds ~3-5 seconds).
+ */
+export function isReady(): boolean {
+  return modelReady;
+}
+
+/**
+ * Initialize the embedding pipeline (singleton).
+ * Lazy: only called on first embed request, not at server startup.
+ * Downloads model on first use, uses cached copy thereafter.
+ */
+async function ensurePipeline(): Promise<void> {
+  if (pipelineInstance) return;
+  if (initPromise) return initPromise;
+
+  initPromise = (async () => {
+    try {
+      logger.info('Initializing embedding model (lazy load)', { model: MODEL_NAME });
+      const { pipeline, env } = await import('@huggingface/transformers');
+
+      // Configure cache location
+      const cacheDir = process.env.TRANSFORMERS_CACHE ?? undefined;
+      if (cacheDir) {
+        env.cacheDir = cacheDir;
+      }
+
+      // Use ONNX backend (default in Transformers.js)
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      pipelineInstance = await (pipeline as any)('feature-extraction', MODEL_NAME, {
+        dtype: 'q8',  // Quantized INT8 for smaller memory footprint
+      });
+
+      modelReady = true;
+      logger.info('Embedding model initialized', { model: MODEL_NAME, dimension: EMBEDDING_DIMENSION });
+    } catch (err) {
+      logger.error('Failed to initialize embedding model', { error: String(err) });
+      initPromise = null;
+      throw err;
+    }
+  })();
+
+  return initPromise;
+}
+
+/**
+ * Choose optimal batch size based on average text length.
+ * Short texts (<100 chars on average) pack more efficiently at batch 64.
+ */
+function chooseBatchSize(texts: string[]): number {
+  if (texts.length === 0) return BASE_BATCH_SIZE;
+  const avgLen = texts.reduce((s, t) => s + t.length, 0) / texts.length;
+  return avgLen < SHORT_TEXT_THRESHOLD ? SHORT_TEXT_BATCH_SIZE : BASE_BATCH_SIZE;
+}
+
+/**
+ * Embed a single text string.
+ * Returns a 384-dimensional float array.
+ */
+export async function embedText(text: string): Promise<number[]> {
+  await ensurePipeline();
+  const result = await pipelineInstance(text, { pooling: 'mean', normalize: true });
+  return Array.from(result.data as Float32Array);
+}
+
+/**
+ * Embed multiple texts in adaptive batches.
+ * Short texts (< SHORT_TEXT_THRESHOLD chars average) use batch size 64.
+ * Long texts use batch size 32.
+ * Returns arrays in the same order as input.
+ * Logs embedding throughput (texts/sec) at info level.
+ */
+export async function embedBatch(texts: string[]): Promise<number[][]> {
+  if (texts.length === 0) return [];
+  await ensurePipeline();
+
+  const batchSize = chooseBatchSize(texts);
+  const results: number[][] = [];
+  const startTime = Date.now();
+
+  for (let i = 0; i < texts.length; i += batchSize) {
+    const batch = texts.slice(i, i + batchSize);
+
+    try {
+      const output = await pipelineInstance(batch, { pooling: 'mean', normalize: true });
+
+      // output.data is a flat Float32Array: [batch_size * dim]
+      const data = output.data as Float32Array;
+      const dim = data.length / batch.length;
+
+      for (let j = 0; j < batch.length; j++) {
+        const vector = Array.from(data.slice(j * dim, (j + 1) * dim));
+        results.push(vector);
+      }
+    } catch (err) {
+      logger.warn('Embedding batch failed, using zero vectors', {
+        batchStart: i,
+        batchSize: batch.length,
+        error: String(err),
+      });
+      // Fallback: push zero vectors for failed batch
+      for (let j = 0; j < batch.length; j++) {
+        results.push(new Array(EMBEDDING_DIMENSION).fill(0));
+      }
+    }
+  }
+
+  const elapsedMs = Date.now() - startTime;
+  const throughput = elapsedMs > 0 ? Math.round((texts.length / elapsedMs) * 1000) : texts.length;
+  logger.info('Embedding batch complete', {
+    totalTexts: texts.length,
+    batchSize,
+    elapsedMs,
+    throughputTextsPerSec: throughput,
+  });
+
+  return results;
+}
+
+/**
+ * Get the embedding dimension.
+ */
+export function getEmbeddingDimension(): number {
+  return EMBEDDING_DIMENSION;
+}
+
+/**
+ * Optionally pre-warm the embedding pipeline in the background.
+ *
+ * Phase 5: This is now a no-op at the call site — initialization is lazy.
+ * The model will load on the first real embed call instead.
+ * This function is kept for backward compatibility; callers that previously
+ * awaited it will continue to work. The model loads asynchronously without
+ * blocking the MCP server startup.
+ */
+export async function warmupEmbedder(): Promise<void> {
+  // Do not eagerly load — lazy initialization is handled by ensurePipeline().
+  // Log a message so callers know the model will warm on first use.
+  logger.info('Embedding pipeline will warm on first use (lazy mode)', { model: MODEL_NAME });
+}