@cue-dev/retrieval-core 0.1.3

package/test/index-prep.concurrency.regression.test.ts

@@ -0,0 +1,142 @@
+import { describe, expect, it } from "vitest";
+import { InMemoryQueryCache } from "@cue-dev/data-plane";
+import { InMemoryIndexStore, RetrievalCore, type EmbeddingProvider } from "../src/index.js";
+
+class TrackingEmbeddingProvider implements EmbeddingProvider {
+  in_flight = 0;
+  max_in_flight = 0;
+
+  reset(): void {
+    this.in_flight = 0;
+    this.max_in_flight = 0;
+  }
+
+  async embed(input: { texts: string[]; purpose: "index" | "query" }): Promise<number[][]> {
+    this.in_flight += 1;
+    this.max_in_flight = Math.max(this.max_in_flight, this.in_flight);
+    try {
+      await new Promise((resolve) => setTimeout(resolve, 15));
+      return input.texts.map(() => [0.1, 0.2, 0.3]);
+    } finally {
+      this.in_flight -= 1;
+    }
+  }
+}
+
+function startTrackingUpsertOrder(store: InMemoryIndexStore): {
+  order: string[];
+  restore: () => void;
+} {
+  const order: string[] = [];
+  const original = store.upsertFile.bind(store);
+  store.upsertFile = async (input) => {
+    order.push(input.repo_path);
+    return original(input);
+  };
+  return {
+    order,
+    restore: () => {
+      store.upsertFile = original;
+    }
+  };
+}
+
+describe("retrieval-core index prep concurrency regression", () => {
+  it("bounds index prep concurrency and uses deterministic code-unit ordering for indexArtifact", async () => {
+    const store = new InMemoryIndexStore();
+    const embeddingProvider = new TrackingEmbeddingProvider();
+    const core = new RetrievalCore(store, new InMemoryQueryCache(), {
+      embeddingProvider,
+      embeddingDescriptor: {
+        provider: "test",
+        model: "tracking",
+        dimensions: 3
+      },
+      indexPrepConcurrency: 2,
+      indexPrepWindowSize: 2
+    });
+
+    await store.upsertWorkspace({
+      workspace_id: "ws-prep-a",
+      tenant_id: "tenant-prep",
+      name: "prep-a",
+      project_root_path: "/workspace/prep-a"
+    });
+
+    const tracker = startTrackingUpsertOrder(store);
+    try {
+      await core.indexArtifact({
+        tenant_id: "tenant-prep",
+        workspace_id: "ws-prep-a",
+        index_version: "idx-prep-a-v1",
+        files: [
+          { path: "src/a.ts", language: "typescript", content: "export const lowerA = 1;\n" },
+          { path: "src/B.ts", language: "typescript", content: "export const upperB = 2;\n" },
+          { path: "src/c.ts", language: "typescript", content: "export const lowerC = 3;\n" }
+        ]
+      });
+    } finally {
+      tracker.restore();
+    }
+
+    expect(embeddingProvider.max_in_flight).toBeGreaterThanOrEqual(2);
+    expect(embeddingProvider.max_in_flight).toBeLessThanOrEqual(2);
+    expect(tracker.order).toEqual(["src/B.ts", "src/a.ts", "src/c.ts"]);
+  });
+
+  it("applies the same bounded deterministic prep path for indexArtifactDelta", async () => {
+    const store = new InMemoryIndexStore();
+    const embeddingProvider = new TrackingEmbeddingProvider();
+    const core = new RetrievalCore(store, new InMemoryQueryCache(), {
+      embeddingProvider,
+      embeddingDescriptor: {
+        provider: "test",
+        model: "tracking",
+        dimensions: 3
+      },
+      indexPrepConcurrency: 2,
+      indexPrepWindowSize: 2
+    });
+
+    await store.upsertWorkspace({
+      workspace_id: "ws-prep-b",
+      tenant_id: "tenant-prep",
+      name: "prep-b",
+      project_root_path: "/workspace/prep-b"
+    });
+
+    await core.indexArtifact({
+      tenant_id: "tenant-prep",
+      workspace_id: "ws-prep-b",
+      index_version: "idx-prep-b-v1",
+      files: [
+        { path: "src/a.ts", language: "typescript", content: "export const lowerA = 1;\n" },
+        { path: "src/B.ts", language: "typescript", content: "export const upperB = 2;\n" },
+        { path: "src/c.ts", language: "typescript", content: "export const lowerC = 3;\n" }
+      ]
+    });
+
+    embeddingProvider.reset();
+    const tracker = startTrackingUpsertOrder(store);
+    try {
+      await core.indexArtifactDelta({
+        tenant_id: "tenant-prep",
+        workspace_id: "ws-prep-b",
+        index_version: "idx-prep-b-v2",
+        base_index_version: "idx-prep-b-v1",
+        upsert_files: [
+          { path: "src/a.ts", language: "typescript", content: "export const lowerA = 10;\n" },
+          { path: "src/Z.ts", language: "typescript", content: "export const upperZ = 20;\n" },
+          { path: "src/B.ts", language: "typescript", content: "export const upperB = 30;\n" }
+        ],
+        deleted_paths: []
+      });
+    } finally {
+      tracker.restore();
+    }
+
+    expect(embeddingProvider.max_in_flight).toBeGreaterThanOrEqual(2);
+    expect(embeddingProvider.max_in_flight).toBeLessThanOrEqual(2);
+    expect(tracker.order).toEqual(["src/B.ts", "src/Z.ts", "src/a.ts"]);
+  });
+});

package/test/integration.test.ts

@@ -0,0 +1,508 @@
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { newDb } from "pg-mem";
+import { InMemoryQueryCache, PostgresIndexRepository } from "@cue-dev/data-plane";
+import { DeterministicEmbeddingProvider, RetrievalCore, RetrievalError, type RerankerProvider } from "../src/index.js";
+
+describe("retrieval-core postgres integration", () => {
+  let repository: PostgresIndexRepository;
+  let close: () => Promise<void>;
+
+  beforeEach(async () => {
+    const db = newDb({ autoCreateForeignKeyIndices: true });
+    const { Pool } = db.adapters.createPg();
+    const pool = new Pool();
+
+    repository = new PostgresIndexRepository(pool, {
+      preferPgVector: false,
+      chunkEmbeddingDimensions: 24
+    });
+    await repository.migrate();
+
+    close = async () => {
+      await pool.end();
+    };
+  });
+
+  afterEach(async () => {
+    await close();
+  });
+
+  async function seedRerankerFixture(): Promise<void> {
+    await repository.upsertWorkspace({
+      workspace_id: "ws-a",
+      tenant_id: "tenant-a",
+      name: "workspace-a",
+      project_root_path: "/workspace/a"
+    });
+
+    const writer = new RetrievalCore(repository, new InMemoryQueryCache());
+    await writer.indexArtifact({
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      index_version: "idx-reranker",
+      files: [
+        {
+          path: "src/primary.ts",
+          language: "typescript",
+          content: "export function selectPrimaryTenant(id: string) { return `primary:${id}`; }"
+        },
+        {
+          path: "src/secondary.ts",
+          language: "typescript",
+          content: "export function selectSecondaryTenant(id: string) { return `secondary:${id}`; }"
+        },
+        {
+          path: "src/helper.ts",
+          language: "typescript",
+          content: "export function tenantLookupHelper(id: string) { return selectPrimaryTenant(id); }"
+        }
+      ]
+    });
+  }
+
+  it("handles indexing lifecycle with added/modified/deleted and secret exclusion", async () => {
+    await repository.upsertWorkspace({
+      workspace_id: "ws-a",
+      tenant_id: "tenant-a",
+      name: "workspace-a",
+      project_root_path: "/workspace/a"
+    });
+
+    const core = new RetrievalCore(repository, new InMemoryQueryCache());
+
+    const reportV1 = await core.indexArtifact({
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      index_version: "idx-v1",
+      files: [
+        {
+          path: "src/app.ts",
+          language: "typescript",
+          content: "export function add(a: number, b: number) { return a + b; }"
+        },
+        {
+          path: "src/utils.ts",
+          language: "typescript",
+          content: "export function clamp(v: number) { return Math.max(0, v); }"
+        },
+        {
+          path: "src/secrets.env",
+          content: "AWS_ACCESS_KEY_ID=AKIAABCDEFGHIJKLMNOP"
+        }
+      ]
+    });
+
+    expect(reportV1.counts.added).toBe(2);
+    expect(reportV1.counts.skipped).toBe(1);
+    expect(reportV1.warnings[0]?.category).toBe("secret_exclusion");
+
+    const readyAfterV1 = await repository.getLatestReadyIndex({
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a"
+    });
+    expect(readyAfterV1).toBeDefined();
+    const metadataV1 = await repository.getIndexMetadata?.({
+      tenant_id: "tenant-a",
+      index_id: readyAfterV1!.index_id
+    });
+    expect(metadataV1?.chunking_strategy).toBe("language_aware");
+    expect(metadataV1?.embedding_provider).toBe("deterministic");
+    expect(metadataV1?.embedding_model).toBe("pseudo-sha256");
+    expect(metadataV1?.embedding_dimensions).toBe(24);
+
+    const reportV2 = await core.indexArtifact({
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      index_version: "idx-v2",
+      files: [
+        {
+          path: "src/app.ts",
+          language: "typescript",
+          content: "export function add(a: number, b: number) { return a + b + 1; }"
+        },
+        {
+          path: "src/new.ts",
+          language: "typescript",
+          content: "export function multiply(a: number, b: number) { return a * b; }"
+        },
+        {
+          path: "src/secrets.env",
+          content: "Bearer abcdefghijklmnopqrstuvwxyz12345"
+        }
+      ]
+    });
+
+    expect(reportV2.counts.modified).toBe(1);
+    expect(reportV2.counts.added).toBe(1);
+    expect(reportV2.counts.deleted).toBe(1);
+    expect(reportV2.counts.skipped).toBe(1);
+
+    const search = await core.searchContext({
+      trace_id: "trc-test",
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      request: {
+        project_root_path: "/workspace/a",
+        query: "multiply function",
+        top_k: 8
+      }
+    });
+
+    expect(search.search_metadata.index_version).toBe("idx-v2");
+    expect(search.results.some((row) => row.path === "src/new.ts")).toBe(true);
+  });
+
+  it("serves retrieval from persisted index across retrieval-core instances", async () => {
+    await repository.upsertWorkspace({
+      workspace_id: "ws-a",
+      tenant_id: "tenant-a",
+      name: "workspace-a",
+      project_root_path: "/workspace/a"
+    });
+
+    const writer = new RetrievalCore(repository, new InMemoryQueryCache());
+    await writer.indexArtifact({
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      index_version: "idx-v1",
+      files: [
+        {
+          path: "src/db.ts",
+          language: "typescript",
+          content: "export const DB_NAME = 'primary';"
+        }
+      ]
+    });
+
+    const reader = new RetrievalCore(repository, new InMemoryQueryCache());
+    const result = await reader.searchContext({
+      trace_id: "trc-read",
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      request: {
+        project_root_path: "/workspace/a",
+        query: "DB_NAME",
+        top_k: 8
+      }
+    });
+
+    expect(result.results.length).toBeGreaterThan(0);
+    expect(result.results[0]?.path).toBe("src/db.ts");
+  });
+
+  it("enforces tenant isolation on retrieval", async () => {
+    await repository.upsertWorkspace({
+      workspace_id: "ws-a",
+      tenant_id: "tenant-a",
+      name: "workspace-a",
+      project_root_path: "/workspace/a"
+    });
+    await repository.upsertWorkspace({
+      workspace_id: "ws-b",
+      tenant_id: "tenant-b",
+      name: "workspace-b",
+      project_root_path: "/workspace/b"
+    });
+
+    const core = new RetrievalCore(repository, new InMemoryQueryCache());
+
+    await core.indexArtifact({
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      index_version: "idx-a1",
+      files: [{ path: "src/a.ts", content: "export const TENANT_A = true;" }]
+    });
+
+    await core.indexArtifact({
+      tenant_id: "tenant-b",
+      workspace_id: "ws-b",
+      index_version: "idx-b1",
+      files: [{ path: "src/b.ts", content: "export const TENANT_B = true;" }]
+    });
+
+    const tenantBSearch = await core.searchContext({
+      trace_id: "trc-b",
+      tenant_id: "tenant-b",
+      workspace_id: "ws-b",
+      request: {
+        project_root_path: "/workspace/b",
+        query: "TENANT_B",
+        top_k: 8
+      }
+    });
+    expect(tenantBSearch.results[0]?.path).toBe("src/b.ts");
+
+    await expect(
+      core.searchContext({
+        trace_id: "trc-cross",
+        tenant_id: "tenant-b",
+        workspace_id: "ws-a",
+        request: {
+          project_root_path: "/workspace/a",
+          query: "TENANT_A",
+          top_k: 8
+        }
+      })
+    ).rejects.toMatchObject({
+      code: "NOT_FOUND"
+    } satisfies Partial<RetrievalError>);
+  });
+
+  it("persists configured embedding metadata on index finalize", async () => {
+    await repository.upsertWorkspace({
+      workspace_id: "ws-a",
+      tenant_id: "tenant-a",
+      name: "workspace-a",
+      project_root_path: "/workspace/a"
+    });
+
+    const core = new RetrievalCore(repository, new InMemoryQueryCache(), {
+      embeddingProvider: new DeterministicEmbeddingProvider({
+        model: "pseudo-sha256-custom",
+        dimensions: 32,
+        version: "v2"
+      })
+    });
+
+    await core.indexArtifact({
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      index_version: "idx-meta",
+      files: [{ path: "src/a.ts", content: "export const A = 1;" }]
+    });
+
+    const latest = await repository.getLatestReadyIndex({
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a"
+    });
+    expect(latest).toBeDefined();
+    const metadata = await repository.getIndexMetadata?.({
+      tenant_id: "tenant-a",
+      index_id: latest!.index_id
+    });
+    expect(metadata?.embedding_provider).toBe("deterministic");
+    expect(metadata?.embedding_model).toBe("pseudo-sha256-custom");
+    expect(metadata?.embedding_dimensions).toBe(32);
+    expect(metadata?.embedding_version).toBe("v2");
+  });
+
+  it("returns INVALID_ARGUMENT when query embedding config is incompatible with latest ready index", async () => {
+    await repository.upsertWorkspace({
+      workspace_id: "ws-a",
+      tenant_id: "tenant-a",
+      name: "workspace-a",
+      project_root_path: "/workspace/a"
+    });
+
+    const writer = new RetrievalCore(repository, new InMemoryQueryCache());
+    await writer.indexArtifact({
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      index_version: "idx-v1",
+      files: [{ path: "src/a.ts", content: "export const A = 1;" }]
+    });
+
+    const request = {
+      trace_id: "trc-mismatch",
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      request: {
+        project_root_path: "/workspace/a",
+        query: "A constant",
+        top_k: 5
+      }
+    } as const;
+
+    const cases = [
+      {
+        name: "provider mismatch",
+        options: {
+          embeddingDescriptor: {
+            provider: "openai_compatible",
+            model: "Qwen/Qwen3-Embedding-4B",
+            dimensions: 24
+          }
+        }
+      },
+      {
+        name: "model mismatch",
+        options: {
+          embeddingDescriptor: {
+            provider: "deterministic",
+            model: "pseudo-sha256-v2",
+            dimensions: 24
+          }
+        }
+      },
+      {
+        name: "dimension mismatch",
+        options: {
+          embeddingDescriptor: {
+            provider: "deterministic",
+            model: "pseudo-sha256",
+            dimensions: 64
+          }
+        }
+      }
+    ] as const;
+
+    for (const testCase of cases) {
+      const reader = new RetrievalCore(repository, new InMemoryQueryCache(), testCase.options);
+      await expect(reader.searchContext(request)).rejects.toMatchObject({
+        code: "INVALID_ARGUMENT"
+      } satisfies Partial<RetrievalError>);
+      await expect(reader.searchContext(request)).rejects.toThrow(/reindex required/i);
+    }
+  });
+
+  it("applies reranker ordering to top-N candidates and avoids cache collisions with deterministic mode", async () => {
+    await seedRerankerFixture();
+
+    const request = {
+      trace_id: "trc-rerank-success",
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      request: {
+        project_root_path: "/workspace/a",
+        query: "select tenant helper",
+        top_k: 3
+      }
+    } as const;
+
+    const sharedCache = new InMemoryQueryCache();
+    const baselineCore = new RetrievalCore(repository, sharedCache);
+    const baseline = await baselineCore.searchContext(request);
+    expect(baseline.results.length).toBeGreaterThanOrEqual(2);
+
+    let rerankCalls = 0;
+    const rerankerProvider: RerankerProvider = {
+      describe: () => ({
+        provider: "test_stub",
+        model: "swap-top-two"
+      }),
+      rerank: async () => {
+        rerankCalls += 1;
+        return [
+          { index: 1, relevance_score: 0.99 },
+          { index: 0, relevance_score: 0.98 }
+        ];
+      }
+    };
+
+    const rerankCore = new RetrievalCore(repository, sharedCache, {
+      rerankerProvider,
+      rerankerTopN: 2
+    });
+    const reranked = await rerankCore.searchContext({
+      ...request,
+      trace_id: "trc-rerank-success-enabled"
+    });
+
+    expect(rerankCalls).toBe(1);
+    expect(reranked.results[0]?.path).toBe(baseline.results[1]?.path);
+  });
+
+  it("falls back to deterministic ranking on reranker timeout", async () => {
+    await seedRerankerFixture();
+
+    const request = {
+      trace_id: "trc-rerank-timeout",
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      request: {
+        project_root_path: "/workspace/a",
+        query: "select tenant helper",
+        top_k: 3
+      }
+    } as const;
+
+    const baselineCore = new RetrievalCore(repository, new InMemoryQueryCache());
+    const baseline = await baselineCore.searchContext(request);
+
+    const timeoutReranker: RerankerProvider = {
+      describe: () => ({
+        provider: "test_stub",
+        model: "timeout"
+      }),
+      rerank: async () => {
+        throw new Error("request timed out after 250ms");
+      }
+    };
+    const rerankCore = new RetrievalCore(repository, new InMemoryQueryCache(), {
+      rerankerProvider: timeoutReranker,
+      rerankerTopN: 3
+    });
+    const timeoutResult = await rerankCore.searchContext({
+      ...request,
+      trace_id: "trc-rerank-timeout-enabled"
+    });
+
+    expect(timeoutResult.results.map((row) => row.path)).toEqual(baseline.results.map((row) => row.path));
+  });
+
+  it("falls back to deterministic ranking on reranker upstream failures", async () => {
+    await seedRerankerFixture();
+
+    const request = {
+      trace_id: "trc-rerank-upstream",
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      request: {
+        project_root_path: "/workspace/a",
+        query: "select tenant helper",
+        top_k: 3
+      }
+    } as const;
+
+    const baselineCore = new RetrievalCore(repository, new InMemoryQueryCache());
+    const baseline = await baselineCore.searchContext(request);
+
+    const failingReranker: RerankerProvider = {
+      describe: () => ({
+        provider: "test_stub",
+        model: "upstream"
+      }),
+      rerank: async () => {
+        throw new Error("HTTP 503 service unavailable");
+      }
+    };
+    const rerankCore = new RetrievalCore(repository, new InMemoryQueryCache(), {
+      rerankerProvider: failingReranker,
+      rerankerTopN: 3
+    });
+    const failedResult = await rerankCore.searchContext({
+      ...request,
+      trace_id: "trc-rerank-upstream-enabled"
+    });
+
+    expect(failedResult.results.map((row) => row.path)).toEqual(baseline.results.map((row) => row.path));
+  });
+
+  it("keeps deterministic baseline behavior when reranker is disabled", async () => {
+    await seedRerankerFixture();
+
+    const request = {
+      trace_id: "trc-rerank-disabled",
+      tenant_id: "tenant-a",
+      workspace_id: "ws-a",
+      request: {
+        project_root_path: "/workspace/a",
+        query: "select tenant helper",
+        top_k: 3
+      }
+    } as const;
+
+    const baselineCore = new RetrievalCore(repository, new InMemoryQueryCache());
+    const baseline = await baselineCore.searchContext(request);
+
+    const disabledRerankCore = new RetrievalCore(repository, new InMemoryQueryCache(), {
+      rerankerTopN: 5
+    });
+    const disabled = await disabledRerankCore.searchContext({
+      ...request,
+      trace_id: "trc-rerank-disabled-explicit"
+    });
+
+    expect(disabled.results.map((row) => row.path)).toEqual(baseline.results.map((row) => row.path));
+  });
+});