@rce-mcp/retrieval-core 0.1.0

package/src/chunking.ts

@@ -0,0 +1,641 @@
+import Parser from "tree-sitter";
+import Go from "tree-sitter-go";
+import JavaScript from "tree-sitter-javascript";
+import Python from "tree-sitter-python";
+import TypeScript from "tree-sitter-typescript";
+
+export type ChunkingStrategy = "language_aware" | "sliding";
+export type ChunkingFallbackReason =
+  | "unsupported_language"
+  | "parser_unavailable"
+  | "parse_error"
+  | "parse_timeout_exceeded"
+  | "empty_language_boundaries";
+
+export interface ChunkingConfig {
+  strategy: ChunkingStrategy;
+  fallback_strategy: "sliding";
+  target_chunk_tokens: number;
+  chunk_overlap_tokens: number;
+  max_chunks_per_file: number;
+  parse_timeout_ms: number;
+  enabled_languages: string[];
+}
+
+export interface ChunkingRawFile {
+  path: string;
+  content: string;
+  language?: string;
+}
+
+export interface ChunkingOutput {
+  start_line: number;
+  end_line: number;
+  snippet: string;
+}
+
+export interface ChunkingResult {
+  chunks: ChunkingOutput[];
+  strategy: ChunkingStrategy;
+  fallback_reason?: ChunkingFallbackReason;
+  parse_latency_ms?: number;
+  language_aware_attempt_latency_ms?: number;
+  fallback_path_latency_ms?: number;
+  language?: string;
+}
+
+export type ParserLanguage = "typescript" | "tsx" | "javascript" | "jsx" | "python" | "go";
+
+type ParserAvailabilityStatus = "available" | "unavailable";
+
+interface ParserAvailability {
+  status: ParserAvailabilityStatus;
+  parser?: Parser;
+  error?: string;
+}
+
+export interface ChunkingParserAvailabilitySnapshotEntry {
+  language: string;
+  status: ParserAvailabilityStatus;
+  error?: string;
+}
+
+const DEFAULT_BOUNDARY_NODE_TYPES: Record<ParserLanguage, Set<string>> = {
+  typescript: new Set([
+    "function_declaration",
+    "generator_function_declaration",
+    "class_declaration",
+    "interface_declaration",
+    "enum_declaration",
+    "type_alias_declaration"
+  ]),
+  tsx: new Set([
+    "function_declaration",
+    "generator_function_declaration",
+    "class_declaration",
+    "interface_declaration",
+    "enum_declaration",
+    "type_alias_declaration"
+  ]),
+  javascript: new Set(["function_declaration", "generator_function_declaration", "class_declaration"]),
+  jsx: new Set(["function_declaration", "generator_function_declaration", "class_declaration"]),
+  python: new Set(["function_definition", "class_definition"]),
+  go: new Set(["function_declaration", "method_declaration", "type_declaration"])
+};
+
+const parserAvailabilityCache = new Map<ParserLanguage, ParserAvailability>();
+const parserInitAttempts = new Map<ParserLanguage, number>();
+const parserLanguageLoaderOverrides = new Map<ParserLanguage, () => Parser.Language>();
+const CANONICAL_TO_PARSER_LANGUAGE: Record<string, ParserLanguage> = {
+  typescript: "typescript",
+  javascript: "javascript",
+  python: "python",
+  go: "go"
+};
+
+function parserLanguageToCanonical(language: ParserLanguage): string {
+  if (language === "tsx") {
+    return "typescript";
+  }
+  if (language === "jsx") {
+    return "javascript";
+  }
+  return language;
+}
+
+function normalizeLanguageAlias(language: string): ParserLanguage | undefined {
+  const normalized = language.trim().toLowerCase();
+  if (normalized === "typescript" || normalized === "ts") {
+    return "typescript";
+  }
+  if (normalized === "tsx") {
+    return "tsx";
+  }
+  if (normalized === "javascript" || normalized === "js") {
+    return "javascript";
+  }
+  if (normalized === "jsx") {
+    return "jsx";
+  }
+  if (normalized === "python" || normalized === "py") {
+    return "python";
+  }
+  if (normalized === "go" || normalized === "golang") {
+    return "go";
+  }
+  return undefined;
+}
+
+function parserLanguageFromPath(path: string): ParserLanguage | undefined {
+  const normalized = path.toLowerCase();
+  if (normalized.endsWith(".tsx")) {
+    return "tsx";
+  }
+  if (normalized.endsWith(".ts")) {
+    return "typescript";
+  }
+  if (normalized.endsWith(".jsx")) {
+    return "jsx";
+  }
+  if (normalized.endsWith(".js") || normalized.endsWith(".mjs") || normalized.endsWith(".cjs")) {
+    return "javascript";
+  }
+  if (normalized.endsWith(".py")) {
+    return "python";
+  }
+  if (normalized.endsWith(".go")) {
+    return "go";
+  }
+  return undefined;
+}
+
+function resolveParserLanguage(file: ChunkingRawFile): ParserLanguage | undefined {
+  const explicit = file.language ? normalizeLanguageAlias(file.language) : undefined;
+  if (explicit) {
+    return explicit;
+  }
+  return parserLanguageFromPath(file.path);
+}
+
+function formatErrorMessage(error: unknown): string {
+  if (error instanceof Error && error.message) {
+    return error.message;
+  }
+  return String(error);
+}
+
+function resolveTreeSitterLanguageHandle(moduleValue: unknown): Parser.Language {
+  if (moduleValue && typeof moduleValue === "object" && "language" in moduleValue) {
+    return (moduleValue as { language: unknown }).language as Parser.Language;
+  }
+  return moduleValue as Parser.Language;
+}
+
+function loadParserLanguage(language: ParserLanguage): Parser.Language {
+  const overrideLoader = parserLanguageLoaderOverrides.get(language);
+  if (overrideLoader) {
+    return overrideLoader();
+  }
+  if (language === "typescript") {
+    return (TypeScript as unknown as { typescript: unknown }).typescript as Parser.Language;
+  }
+  if (language === "tsx") {
+    return (TypeScript as unknown as { tsx: unknown }).tsx as Parser.Language;
+  }
+  if (language === "javascript" || language === "jsx") {
+    return resolveTreeSitterLanguageHandle(JavaScript);
+  }
+  if (language === "python") {
+    return resolveTreeSitterLanguageHandle(Python);
+  }
+  return resolveTreeSitterLanguageHandle(Go);
+}
+
+function getParserAvailability(language: ParserLanguage): ParserAvailability {
+  const cached = parserAvailabilityCache.get(language);
+  if (cached) {
+    return cached;
+  }
+  parserInitAttempts.set(language, (parserInitAttempts.get(language) ?? 0) + 1);
+  try {
+    const parser = new Parser();
+    parser.setLanguage(loadParserLanguage(language));
+    const availability: ParserAvailability = {
+      status: "available",
+      parser
+    };
+    parserAvailabilityCache.set(language, availability);
+    return availability;
+  } catch (error) {
+    const availability: ParserAvailability = {
+      status: "unavailable",
+      error: formatErrorMessage(error)
+    };
+    parserAvailabilityCache.set(language, availability);
+    return availability;
+  }
+}
+
+function getParser(language: ParserLanguage): Parser | undefined {
+  const availability = getParserAvailability(language);
+  if (availability.status !== "available") {
+    return undefined;
+  }
+  return availability.parser;
+}
+
+function toChunkingParserLanguage(canonicalLanguage: string): ParserLanguage | undefined {
+  return CANONICAL_TO_PARSER_LANGUAGE[canonicalLanguage];
+}
+
+export function getChunkingParserAvailabilitySnapshot(input?: {
+  enabled_languages?: string[];
+}): ChunkingParserAvailabilitySnapshotEntry[] {
+  const canonicalLanguages = new Set<string>();
+  if (input?.enabled_languages && input.enabled_languages.length > 0) {
+    for (const language of input.enabled_languages) {
+      const parsed = normalizeLanguageAlias(language);
+      if (!parsed) {
+        continue;
+      }
+      const canonical = parserLanguageToCanonical(parsed);
+      if (toChunkingParserLanguage(canonical)) {
+        canonicalLanguages.add(canonical);
+      }
+    }
+  }
+  if (canonicalLanguages.size === 0) {
+    for (const canonical of Object.keys(CANONICAL_TO_PARSER_LANGUAGE)) {
+      canonicalLanguages.add(canonical);
+    }
+  }
+
+  return [...canonicalLanguages]
+    .sort((a, b) => a.localeCompare(b))
+    .map((canonicalLanguage) => {
+      const parserLanguage = toChunkingParserLanguage(canonicalLanguage);
+      if (!parserLanguage) {
+        return {
+          language: canonicalLanguage,
+          status: "unavailable" as const,
+          error: "no parser mapping for language"
+        };
+      }
+      const availability = getParserAvailability(parserLanguage);
+      return {
+        language: canonicalLanguage,
+        status: availability.status,
+        ...(availability.error ? { error: availability.error } : {})
+      };
+    });
+}
+
+export function __resetChunkingParserStateForTests(): void {
+  parserAvailabilityCache.clear();
+  parserInitAttempts.clear();
+  parserLanguageLoaderOverrides.clear();
+}
+
+export function __setChunkingParserLanguageLoaderForTests(
+  language: ParserLanguage,
+  loader: (() => Parser.Language) | undefined
+): void {
+  if (loader) {
+    parserLanguageLoaderOverrides.set(language, loader);
+  } else {
+    parserLanguageLoaderOverrides.delete(language);
+  }
+  parserAvailabilityCache.delete(language);
+}
+
+export function __getChunkingParserInitAttemptsForTests(): Record<string, number> {
+  return Object.fromEntries(parserInitAttempts.entries());
+}
+
+function toInclusiveEndRow(node: Parser.SyntaxNode): number {
+  const end = node.endPosition;
+  if (end.column === 0 && end.row > node.startPosition.row) {
+    return end.row - 1;
+  }
+  return end.row;
+}
+
+function trimLineRange(lines: string[], startRow: number, endRow: number): { start: number; end: number } | undefined {
+  let start = startRow;
+  let end = endRow;
+  while (start <= end && (lines[start] ?? "").trim().length === 0) {
+    start += 1;
+  }
+  while (end >= start && (lines[end] ?? "").trim().length === 0) {
+    end -= 1;
+  }
+  if (end < start) {
+    return undefined;
+  }
+  return { start, end };
+}
+
+function splitRangeWithBudget(input: {
+  lines: string[];
+  startRow: number;
+  endRow: number;
+  tokenize: (text: string) => string[];
+  targetChunkTokens: number;
+  overlapTokens: number;
+  maxChunks: number;
+}): Array<{ startRow: number; endRow: number }> {
+  const segments: Array<{ startRow: number; endRow: number }> = [];
+  let start = input.startRow;
+
+  while (start <= input.endRow && segments.length < input.maxChunks) {
+    let tokens = 0;
+    let end = start;
+    while (end <= input.endRow) {
+      tokens += input.tokenize(input.lines[end] ?? "").length;
+      if (tokens >= input.targetChunkTokens) {
+        break;
+      }
+      end += 1;
+    }
+
+    const safeEnd = Math.min(end, input.endRow);
+    if (safeEnd >= start) {
+      segments.push({ startRow: start, endRow: safeEnd });
+    }
+
+    if (safeEnd >= input.endRow) {
+      break;
+    }
+    const rewind = Math.max(1, Math.floor(input.overlapTokens / 4));
+    start = Math.max(start + 1, safeEnd - rewind + 1);
+  }
+
+  return segments;
+}
+
+function buildSlidingChunks(input: {
+  lines: string[];
+  tokenize: (text: string) => string[];
+  targetChunkTokens: number;
+  overlapTokens: number;
+  maxChunks: number;
+}): ChunkingOutput[] {
+  const rawSegments = splitRangeWithBudget({
+    lines: input.lines,
+    startRow: 0,
+    endRow: Math.max(0, input.lines.length - 1),
+    tokenize: input.tokenize,
+    targetChunkTokens: input.targetChunkTokens,
+    overlapTokens: input.overlapTokens,
+    maxChunks: input.maxChunks
+  });
+  const chunks: ChunkingOutput[] = [];
+  for (const segment of rawSegments) {
+    const trimmed = trimLineRange(input.lines, segment.startRow, segment.endRow);
+    if (!trimmed) {
+      continue;
+    }
+    chunks.push({
+      start_line: trimmed.start + 1,
+      end_line: trimmed.end + 1,
+      snippet: input.lines.slice(trimmed.start, trimmed.end + 1).join("\n")
+    });
+  }
+  return chunks;
+}
+
+function hasBoundaryAncestor(node: Parser.SyntaxNode, boundaryTypes: Set<string>): boolean {
+  let current = node.parent;
+  while (current) {
+    if (boundaryTypes.has(current.type)) {
+      return true;
+    }
+    current = current.parent;
+  }
+  return false;
+}
+
+function buildLanguageAwareChunks(input: {
+  file: ChunkingRawFile;
+  lines: string[];
+  parserLanguage: ParserLanguage;
+  config: ChunkingConfig;
+  tokenize: (text: string) => string[];
+}): ChunkingResult {
+  const languageAwareAttemptStart = Date.now();
+  const parser = getParser(input.parserLanguage);
+  if (!parser) {
+    const fallbackStart = Date.now();
+    const chunks = buildSlidingChunks({
+      lines: input.lines,
+      tokenize: input.tokenize,
+      targetChunkTokens: input.config.target_chunk_tokens,
+      overlapTokens: input.config.chunk_overlap_tokens,
+      maxChunks: input.config.max_chunks_per_file
+    });
+    return {
+      chunks,
+      strategy: "sliding",
+      fallback_reason: "parser_unavailable",
+      language_aware_attempt_latency_ms: Date.now() - languageAwareAttemptStart,
+      fallback_path_latency_ms: Date.now() - fallbackStart,
+      language: parserLanguageToCanonical(input.parserLanguage)
+    };
+  }
+
+  try {
+    parser.setTimeoutMicros(input.config.parse_timeout_ms * 1_000);
+    const parseStart = Date.now();
+    const tree = parser.parse(input.file.content);
+    const parseLatencyMs = Date.now() - parseStart;
+    if (parseLatencyMs > input.config.parse_timeout_ms) {
+      const fallbackStart = Date.now();
+      const chunks = buildSlidingChunks({
+        lines: input.lines,
+        tokenize: input.tokenize,
+        targetChunkTokens: input.config.target_chunk_tokens,
+        overlapTokens: input.config.chunk_overlap_tokens,
+        maxChunks: input.config.max_chunks_per_file
+      });
+      return {
+        chunks,
+        strategy: "sliding",
+        fallback_reason: "parse_timeout_exceeded",
+        parse_latency_ms: parseLatencyMs,
+        language_aware_attempt_latency_ms: Date.now() - languageAwareAttemptStart,
+        fallback_path_latency_ms: Date.now() - fallbackStart,
+        language: parserLanguageToCanonical(input.parserLanguage)
+      };
+    }
+
+    const root = (tree as Parser.Tree | null)?.rootNode;
+    if (!root) {
+      const fallbackStart = Date.now();
+      const chunks = buildSlidingChunks({
+        lines: input.lines,
+        tokenize: input.tokenize,
+        targetChunkTokens: input.config.target_chunk_tokens,
+        overlapTokens: input.config.chunk_overlap_tokens,
+        maxChunks: input.config.max_chunks_per_file
+      });
+      return {
+        chunks,
+        strategy: "sliding",
+        fallback_reason: "parse_error",
+        parse_latency_ms: parseLatencyMs,
+        language_aware_attempt_latency_ms: Date.now() - languageAwareAttemptStart,
+        fallback_path_latency_ms: Date.now() - fallbackStart,
+        language: parserLanguageToCanonical(input.parserLanguage)
+      };
+    }
+
+    const boundaryTypes = DEFAULT_BOUNDARY_NODE_TYPES[input.parserLanguage];
+    const candidates = root.descendantsOfType([...boundaryTypes]);
+    const boundaryNodes = candidates
+      .filter((node) => !hasBoundaryAncestor(node, boundaryTypes))
+      .sort((a, b) => a.startPosition.row - b.startPosition.row || a.startPosition.column - b.startPosition.column);
+
+    if (boundaryNodes.length === 0) {
+      const fallbackStart = Date.now();
+      const chunks = buildSlidingChunks({
+        lines: input.lines,
+        tokenize: input.tokenize,
+        targetChunkTokens: input.config.target_chunk_tokens,
+        overlapTokens: input.config.chunk_overlap_tokens,
+        maxChunks: input.config.max_chunks_per_file
+      });
+      return {
+        chunks,
+        strategy: "sliding",
+        fallback_reason: "empty_language_boundaries",
+        parse_latency_ms: parseLatencyMs,
+        language_aware_attempt_latency_ms: Date.now() - languageAwareAttemptStart,
+        fallback_path_latency_ms: Date.now() - fallbackStart,
+        language: parserLanguageToCanonical(input.parserLanguage)
+      };
+    }
+
+    const segments: Array<{ startRow: number; endRow: number }> = [];
+    let cursor = 0;
+    const lastRow = Math.max(0, input.lines.length - 1);
+    for (const node of boundaryNodes) {
+      const startRow = Math.max(0, Math.min(lastRow, node.startPosition.row));
+      const endRow = Math.max(startRow, Math.min(lastRow, toInclusiveEndRow(node)));
+      if (startRow > cursor) {
+        segments.push({ startRow: cursor, endRow: startRow - 1 });
+      }
+      segments.push({ startRow, endRow });
+      cursor = endRow + 1;
+      if (cursor > lastRow) {
+        break;
+      }
+    }
+    if (cursor <= lastRow) {
+      segments.push({ startRow: cursor, endRow: lastRow });
+    }
+
+    const chunks: ChunkingOutput[] = [];
+    for (const segment of segments) {
+      if (segment.endRow < segment.startRow || chunks.length >= input.config.max_chunks_per_file) {
+        continue;
+      }
+      const pieces = splitRangeWithBudget({
+        lines: input.lines,
+        startRow: segment.startRow,
+        endRow: segment.endRow,
+        tokenize: input.tokenize,
+        targetChunkTokens: input.config.target_chunk_tokens,
+        overlapTokens: input.config.chunk_overlap_tokens,
+        maxChunks: input.config.max_chunks_per_file - chunks.length
+      });
+      for (const piece of pieces) {
+        const trimmed = trimLineRange(input.lines, piece.startRow, piece.endRow);
+        if (!trimmed) {
+          continue;
+        }
+        chunks.push({
+          start_line: trimmed.start + 1,
+          end_line: trimmed.end + 1,
+          snippet: input.lines.slice(trimmed.start, trimmed.end + 1).join("\n")
+        });
+        if (chunks.length >= input.config.max_chunks_per_file) {
+          break;
+        }
+      }
+    }
+
+    if (chunks.length === 0) {
+      const fallbackStart = Date.now();
+      const slidingChunks = buildSlidingChunks({
+        lines: input.lines,
+        tokenize: input.tokenize,
+        targetChunkTokens: input.config.target_chunk_tokens,
+        overlapTokens: input.config.chunk_overlap_tokens,
+        maxChunks: input.config.max_chunks_per_file
+      });
+      return {
+        chunks: slidingChunks,
+        strategy: "sliding",
+        fallback_reason: "empty_language_boundaries",
+        parse_latency_ms: parseLatencyMs,
+        language_aware_attempt_latency_ms: Date.now() - languageAwareAttemptStart,
+        fallback_path_latency_ms: Date.now() - fallbackStart,
+        language: parserLanguageToCanonical(input.parserLanguage)
+      };
+    }
+
+    return {
+      chunks,
+      strategy: "language_aware",
+      parse_latency_ms: parseLatencyMs,
+      language_aware_attempt_latency_ms: Date.now() - languageAwareAttemptStart,
+      language: parserLanguageToCanonical(input.parserLanguage)
+    };
+  } catch {
+    const fallbackStart = Date.now();
+    const chunks = buildSlidingChunks({
+      lines: input.lines,
+      tokenize: input.tokenize,
+      targetChunkTokens: input.config.target_chunk_tokens,
+      overlapTokens: input.config.chunk_overlap_tokens,
+      maxChunks: input.config.max_chunks_per_file
+    });
+    return {
+      chunks,
+      strategy: "sliding",
+      fallback_reason: "parse_error",
+      language_aware_attempt_latency_ms: Date.now() - languageAwareAttemptStart,
+      fallback_path_latency_ms: Date.now() - fallbackStart,
+      language: parserLanguageToCanonical(input.parserLanguage)
+    };
+  }
+}
+
+export function buildChunksForFile(input: {
+  file: ChunkingRawFile;
+  config: ChunkingConfig;
+  tokenize: (text: string) => string[];
+}): ChunkingResult {
+  const lines = input.file.content.split("\n");
+  const language = resolveParserLanguage(input.file);
+  const enabledLanguageSet = new Set(input.config.enabled_languages.map((value) => value.trim().toLowerCase()));
+
+  if (input.config.strategy === "sliding") {
+    return {
+      chunks: buildSlidingChunks({
+        lines,
+        tokenize: input.tokenize,
+        targetChunkTokens: input.config.target_chunk_tokens,
+        overlapTokens: input.config.chunk_overlap_tokens,
+        maxChunks: input.config.max_chunks_per_file
+      }),
+      strategy: "sliding",
+      language: language ? parserLanguageToCanonical(language) : undefined
+    };
+  }
+
+  if (!language || !enabledLanguageSet.has(parserLanguageToCanonical(language))) {
+    const fallbackStart = Date.now();
+    const chunks = buildSlidingChunks({
+      lines,
+      tokenize: input.tokenize,
+      targetChunkTokens: input.config.target_chunk_tokens,
+      overlapTokens: input.config.chunk_overlap_tokens,
+      maxChunks: input.config.max_chunks_per_file
+    });
+    return {
+      chunks,
+      strategy: "sliding",
+      fallback_reason: "unsupported_language",
+      fallback_path_latency_ms: Date.now() - fallbackStart
+    };
+  }
+
+  return buildLanguageAwareChunks({
+    file: input.file,
+    lines,
+    parserLanguage: language,
+    config: input.config,
+    tokenize: input.tokenize
+  });
+}