GitNexus/gitnexus/test/unit/embedding-chunking.test.ts

/**
 * Integration test: Embedding chunking pipeline
 *
 * Tests the chunking + text generation pipeline together.
 */
import { describe, it, expect, vi } from 'vitest';
import { characterChunk } from '../../src/core/embeddings/character-chunk.js';
import { generateEmbeddingText } from '../../src/core/embeddings/text-generator.js';
import type { EmbeddableNode } from '../../src/core/embeddings/types.js';

const { createParserForLanguage } = vi.hoisted(() => ({
  createParserForLanguage: vi.fn(),
}));

vi.mock('../../src/core/tree-sitter/parser-loader.js', () => ({
  createParserForLanguage,
  isLanguageAvailable: vi.fn().mockReturnValue(true),
  resolveLanguageKey: vi.fn((language: string) => language),
}));

const { getLanguageFromFilename } = vi.hoisted(() => ({
  getLanguageFromFilename: vi.fn().mockReturnValue('typescript'),
}));

// Partial mock: `ast-utils` now resolves the LanguageProvider registry to apply
// `preprocessSource`, and that graph needs the real shared exports (#2771).
vi.mock('gitnexus-shared', async (importOriginal) => ({
  ...(await importOriginal<typeof import('gitnexus-shared')>()),
  getLanguageFromFilename,
}));

import { chunkNode } from '../../src/core/embeddings/chunker.js';

const CLASS_PREV_TAIL_SAMPLE = 30;
const STRUCT_PREV_TAIL_SAMPLE = 20;

type FakeNode = {
  type: string;
  startIndex: number;
  endIndex: number;
  namedChildCount: number;
  namedChild: (index: number) => FakeNode | null;
  childForFieldName?: (name: string) => FakeNode | null;
};

const makeFakeNode = (
  type: string,
  startIndex: number,
  endIndex: number,
  children: FakeNode[] = [],
  fields: Record<string, FakeNode> = {},
): FakeNode => ({
  type,
  startIndex,
  endIndex,
  namedChildCount: children.length,
  namedChild: (index: number) => children[index] ?? null,
  childForFieldName: (name: string) => fields[name] ?? null,
});

const makeDeclarationTree = (
  nodeType: string,
  bodyType: string,
  content: string,
  members: Array<{ text: string; type: string }>,
) => {
  let searchFrom = 0;
  const memberNodes = members.map((member) => {
    const startIndex = content.indexOf(member.text, searchFrom);
    if (startIndex < 0) {
      throw new Error(`Unable to locate declaration member text: ${member.text}`);
    }
    searchFrom = startIndex + member.text.length;
    return makeFakeNode(member.type, startIndex, startIndex + member.text.length);
  });

  const bodyStart = content.indexOf('{');
  const bodyEnd = content.lastIndexOf('}') + 1;
  const bodyNode = makeFakeNode(bodyType, bodyStart, bodyEnd, memberNodes);
  const declNode = makeFakeNode(nodeType, 0, bodyEnd, [bodyNode], { body: bodyNode });
  return {
    rootNode: makeFakeNode('program', 0, content.length, [declNode]),
  };
};

describe('embedding-chunking integration', () => {
  const makeNode = (overrides: Partial<EmbeddableNode>): EmbeddableNode => ({
    id: 'Function:src/test.ts:test',
    name: 'test',
    label: 'Function',
    filePath: 'src/test.ts',
    content: '',
    startLine: 1,
    endLine: 10,
    ...overrides,
  });

  it('short function produces single chunk with metadata', () => {
    const node = makeNode({
      content: 'function hello() { return "world"; }',
      isExported: true,
      repoName: 'my-project',
      serverName: 'my-service',
    });

    const chunks = characterChunk(node.content, 1, 3, 1200, 120);
    expect(chunks).toHaveLength(1);

    const text = generateEmbeddingText(node, chunks[0].text);
    expect(text).toContain('Function: test');
    expect(text).toContain('function hello()');
    // #2333: verbose metadata is no longer part of embedding text.
    expect(text).not.toContain('Repo: my-project');
    expect(text).not.toContain('Server: my-service');
    expect(text).not.toContain('Export: true');
  });

  it('long function produces multiple chunks', () => {
    const longContent = Array.from({ length: 100 }, (_, i) => `  const line${i} = ${i};`).join(
      '\n',
    );
    const node = makeNode({
      content: `function longFn() {\n${longContent}\n}`,
      startLine: 1,
      endLine: 102,
    });

    const chunks = characterChunk(node.content, 1, 102, 1200, 120);
    expect(chunks.length).toBeGreaterThan(1);
    expect(chunks[0].chunkIndex).toBe(0);
  });

  it('short labels (TypeAlias) skip chunking and embed directly', () => {
    const node = makeNode({
      label: 'TypeAlias',
      name: 'Result',
      content: 'type Result<T> = Success<T> | Error;',
    });

    const chunks = characterChunk(node.content, 1, 1, 1200, 120);
    expect(chunks).toHaveLength(1);

    const text = generateEmbeddingText(node, chunks[0].text);
    expect(text).toContain('TypeAlias: Result');
    expect(text).toContain('type Result<T> = Success<T> | Error;');
  });

  it('long enum uses character fallback', () => {
    const enumContent = Array.from(
      { length: 200 },
      (_, i) => `  Value${i} = "${'x'.repeat(20)}${i}",`,
    ).join('\n');
    const node = makeNode({
      label: 'Enum',
      name: 'LargeEnum',
      content: `enum LargeEnum {\n${enumContent}\n}`,
      startLine: 1,
      endLine: 202,
    });

    const chunks = characterChunk(node.content, 1, 202, 1200, 120);
    expect(chunks.length).toBeGreaterThan(1);
  });

  it('long class produces member-aware chunks with structural metadata', async () => {
    const node = makeNode({
      label: 'Class',
      name: 'Parser',
      methodNames: ['parseJSON', 'validate'],
      fieldNames: ['options', 'cache'],
      content: `class Parser {
  options: ParserOptions;
  cache: Map<string, any>;
  parseJSON(text: string) { return JSON.parse(text); }
  validate() { return true; }
}`,
      startLine: 20,
      endLine: 25,
    });
    createParserForLanguage.mockResolvedValue({
      parse: vi.fn().mockReturnValue(
        makeDeclarationTree('class_declaration', 'class_body', node.content, [
          { text: 'options: ParserOptions;', type: 'field_definition' },
          { text: 'cache: Map<string, any>;', type: 'field_definition' },
          {
            text: 'parseJSON(text: string) { return JSON.parse(text); }',
            type: 'method_definition',
          },
          { text: 'validate() { return true; }', type: 'method_definition' },
        ]),
      ),
    });

    const chunks = await chunkNode(node.label, node.content, node.filePath, 20, 25, 90, 0);
    expect(chunks).toHaveLength(2);

    const secondText = generateEmbeddingText(
      node,
      chunks[1].text,
      {},
      chunks[1].chunkIndex,
      chunks[0].text.slice(-CLASS_PREV_TAIL_SAMPLE),
    );
    expect(secondText).toContain('Class: Parser');
    expect(secondText).toContain('Container: class Parser {');
    expect(secondText).toContain('[preceding context]: ...');
    expect(secondText).not.toContain('Methods: parseJSON, validate');
    expect(secondText).not.toContain('Properties: options, cache');
    expect(secondText).toContain('parseJSON(text: string)');
  });

  it('interface chunks retain structural metadata and signatures', async () => {
    const node = makeNode({
      label: 'Interface',
      name: 'Handler',
      methodNames: ['handle', 'validate'],
      fieldNames: ['name'],
      content: `interface Handler {
  handle(event: Event): void;
  validate(input: string): boolean;
  readonly name: string;
}`,
      startLine: 30,
      endLine: 34,
    });
    createParserForLanguage.mockResolvedValue({
      parse: vi.fn().mockReturnValue(
        makeDeclarationTree('interface_declaration', 'object_type', node.content, [
          { text: 'handle(event: Event): void;', type: 'method_definition' },
          { text: 'validate(input: string): boolean;', type: 'method_definition' },
          { text: 'readonly name: string;', type: 'property_signature' },
        ]),
      ),
    });

    const chunks = await chunkNode(node.label, node.content, node.filePath, 30, 34, 500, 0);
    expect(chunks).toHaveLength(1);

    const text = generateEmbeddingText(node, chunks[0].text);
    expect(text).toContain('Interface: Handler');
    expect(text).toContain('Methods: handle, validate');
    expect(text).toContain('Container: interface Handler {');
    expect(text).toContain('readonly name: string;');
  });

  it('struct chunks retain structural container context', async () => {
    getLanguageFromFilename.mockReturnValue('rust');
    const node = makeNode({
      label: 'Struct',
      name: 'User',
      fieldNames: ['name', 'email', 'age', 'address'],
      content: `struct User {
  name: String,
  email: String,
  age: u32,
  address: String,
}`,
      startLine: 40,
      endLine: 45,
      filePath: 'src/user.rs',
    });
    createParserForLanguage.mockResolvedValue({
      parse: vi.fn().mockReturnValue(
        makeDeclarationTree('struct_item', 'declaration_list', node.content, [
          { text: 'name: String,', type: 'field_definition' },
          { text: 'email: String,', type: 'field_definition' },
          { text: 'age: u32,', type: 'field_definition' },
          { text: 'address: String,', type: 'field_definition' },
        ]),
      ),
    });

    const chunks = await chunkNode(node.label, node.content, node.filePath, 40, 45, 45, 0);
    expect(chunks).toHaveLength(2);

    const secondText = generateEmbeddingText(
      node,
      chunks[1].text,
      {},
      chunks[1].chunkIndex,
      chunks[0].text.slice(-STRUCT_PREV_TAIL_SAMPLE),
    );
    expect(secondText).toContain('Struct: User');
    expect(secondText).toContain('Container: struct User {');
    expect(secondText).not.toContain('Properties: name, email, age, address');
    expect(secondText).toContain('age: u32,');
  });

  it('header is present in every chunk', () => {
    const longContent = 'x'.repeat(3000);
    const node = makeNode({
      content: longContent,
      repoName: 'test-repo',
    });

    const chunks = characterChunk(node.content, 1, 100, 1200, 120);
    expect(chunks.length).toBeGreaterThan(1);

    for (const chunk of chunks) {
      const text = generateEmbeddingText(node, chunk.text);
      // The compact header (name, + description when present) repeats on every
      // chunk so each chunk keeps its identity; #2333 dropped the metadata lines.
      expect(text).toContain('Function: test');
      expect(text).not.toContain('Repo: test-repo');
      expect(text).not.toContain('Path: src/test.ts');
    }
  });
});