@axiom-lattice/pg-stores 5.0.0 → 5.0.1

package/src/__tests__/PostgreSQLEvalStore.integration.test.ts

@@ -0,0 +1,575 @@
+import { randomUUID } from "node:crypto";
+import { Pool } from "pg";
+
+import { PostgreSQLEvalStore } from "../stores/PostgreSQLEvalStore";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const describeIfDatabase = TEST_DATABASE_URL ? describe : describe.skip;
+
+describeIfDatabase("PostgreSQLEvalStore integration", () => {
+  let store: PostgreSQLEvalStore;
+  let pool: Pool;
+  const tenants = new Set<string>();
+
+  beforeAll(async () => {
+    store = new PostgreSQLEvalStore({
+      poolConfig: TEST_DATABASE_URL!,
+      autoMigrate: true,
+    });
+    await store.initialize();
+    pool = new Pool({ connectionString: TEST_DATABASE_URL });
+  });
+
+  afterEach(async () => {
+    if (tenants.size === 0) {
+      return;
+    }
+    const ids = [...tenants];
+    await pool.query(
+      `DELETE FROM lattice_eval_run_results
+       WHERE run_id IN (SELECT id FROM lattice_eval_runs WHERE tenant_id = ANY($1::text[]))`,
+      [ids],
+    );
+    await pool.query("DELETE FROM lattice_eval_runs WHERE tenant_id = ANY($1::text[])", [ids]);
+    await pool.query("DELETE FROM lattice_eval_cases WHERE tenant_id = ANY($1::text[])", [ids]);
+    await pool.query("DELETE FROM lattice_eval_suites WHERE tenant_id = ANY($1::text[])", [ids]);
+    await pool.query("DELETE FROM lattice_eval_projects WHERE tenant_id = ANY($1::text[])", [ids]);
+    tenants.clear();
+  });
+
+  afterAll(async () => {
+    await pool.end();
+    await store.dispose();
+  });
+
+  function tenantId(): string {
+    const id = `eval-${randomUUID()}`;
+    tenants.add(id);
+    return id;
+  }
+
+  async function createProject(tenant: string, name = "Project") {
+    return store.createProject(tenant, randomUUID(), {
+      name,
+      description: "A test project",
+      version: "1.0.0",
+      judgeModelConfig: { provider: "openai", model: "gpt-4o" },
+      targetServerConfig: { baseUrl: "https://agent.example" },
+      concurrency: 2,
+      reportConfig: { format: "markdown" },
+    });
+  }
+
+  async function createSuite(tenant: string, projectId: string, name = "Suite") {
+    return store.createSuite(tenant, projectId, randomUUID(), { name });
+  }
+
+  async function createCase(tenant: string, suiteId: string) {
+    return store.createCase(tenant, suiteId, randomUUID(), {
+      inputMessage: "Approve the payment",
+      inputFiles: { receipt: "file-1" },
+      steps: [{ agent_id: "executor-1", override_message: "go" }],
+      outputType: "file_content",
+      contentAssertion: "payment executed",
+      rubrics: [{ name: "accuracy", weight: 0.5, description: "Correct result" }],
+      interruptPolicy: { mode: "canned-response", value: "同意" },
+    });
+  }
+
+  async function createRun(tenant: string, projectId: string) {
+    return store.createRun(tenant, projectId, randomUUID(), {
+      totalCases: 4,
+      concurrency: 2,
+      holdout: true,
+      envProjectId: "env-project",
+      envWorkspaceId: "env-workspace",
+      taskId: "task-1",
+    });
+  }
+
+  async function createResult(tenant: string, runId: string) {
+    return store.createRunResult(tenant, runId, randomUUID(), {
+      suiteName: "main",
+      pass: true,
+      score: 0.9,
+      summary: "looks good",
+      dimensionResults: [{ name: "quality", score: 0.9, reason: "clear" }],
+      durationMs: 120,
+      messages: [{ role: "assistant", content: "done", id: "m1" }],
+      logs: [{ timestamp: "2026-01-01T00:00:00.000Z", level: "info", message: "ran" }],
+      interrupted: false,
+    });
+  }
+
+  describe("projects", () => {
+    it("creates a project and reads it back by id", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant, "Support bot");
+
+      expect(project.tenantId).toBe(tenant);
+      expect(project.name).toBe("Support bot");
+      expect(project.description).toBe("A test project");
+      expect(project.version).toBe("1.0.0");
+      expect(project.judgeModelConfig).toEqual({ provider: "openai", model: "gpt-4o" });
+      expect(project.targetServerConfig).toEqual({ baseUrl: "https://agent.example" });
+      expect(project.concurrency).toBe(2);
+      expect(project.reportConfig).toEqual({ format: "markdown" });
+      expect(project.createdAt).toBeInstanceOf(Date);
+      expect(project.updatedAt).toBeInstanceOf(Date);
+
+      const read = await store.getProjectById(tenant, project.id);
+      expect(read).toMatchObject({
+        id: project.id,
+        tenantId: tenant,
+        name: "Support bot",
+        concurrency: 2,
+      });
+    });
+
+    it("lists only the requested tenant's projects", async () => {
+      const tenantA = tenantId();
+      const tenantB = tenantId();
+      const a1 = await createProject(tenantA, "A1");
+      const a2 = await createProject(tenantA, "A2");
+      const b1 = await createProject(tenantB, "B1");
+
+      const listedA = await store.getProjectsByTenant(tenantA);
+      expect(listedA.map((p) => p.id).sort()).toEqual([a1.id, a2.id].sort());
+      expect(listedA.every((p) => p.tenantId === tenantA)).toBe(true);
+
+      const listedB = await store.getProjectsByTenant(tenantB);
+      expect(listedB.map((p) => p.id)).toEqual([b1.id]);
+    });
+
+    it("updates only the provided project fields", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant, "Before");
+
+      const updated = await store.updateProject(tenant, project.id, {
+        name: "After",
+        concurrency: 5,
+      });
+
+      expect(updated).toMatchObject({
+        id: project.id,
+        name: "After",
+        concurrency: 5,
+        description: "A test project",
+        version: "1.0.0",
+      });
+    });
+
+    it("returns null / false for an absent project id", async () => {
+      const tenant = tenantId();
+      const missing = randomUUID();
+
+      expect(await store.getProjectById(tenant, missing)).toBeNull();
+      expect(await store.updateProject(tenant, missing, { name: "nope" })).toBeNull();
+      expect(await store.deleteProject(tenant, missing)).toBe(false);
+      expect(await store.getProjectReport(tenant, missing)).toBeNull();
+    });
+
+    it("deletes a project", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+
+      expect(await store.deleteProject(tenant, project.id)).toBe(true);
+      expect(await store.getProjectById(tenant, project.id)).toBeNull();
+      expect(await store.deleteProject(tenant, project.id)).toBe(false);
+    });
+
+    it("defaults an omitted judgeModelConfig to an empty object", async () => {
+      const tenant = tenantId();
+      const project = await store.createProject(tenant, randomUUID(), {
+        name: "No judge configured",
+        targetServerConfig: { baseUrl: "https://agent.example" },
+      });
+
+      expect(project.judgeModelConfig).toEqual({});
+
+      const read = await store.getProjectById(tenant, project.id);
+      expect(read!.judgeModelConfig).toEqual({});
+    });
+  });
+
+  describe("suites", () => {
+    it("creates suites scoped to a project and lists them with case counts", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const otherProject = await createProject(tenant, "Other");
+      const suiteA = await createSuite(tenant, project.id, "Smoke");
+      const suiteB = await createSuite(tenant, project.id, "Regression");
+      await createSuite(tenant, otherProject.id, "Unrelated");
+      await createCase(tenant, suiteA.id);
+
+      const listed = await store.getSuitesByProject(tenant, project.id);
+      expect(listed.map((s) => s.id).sort()).toEqual([suiteA.id, suiteB.id].sort());
+      expect(listed.find((s) => s.id === suiteA.id)?.caseCount).toBe(1);
+      expect(listed.find((s) => s.id === suiteB.id)?.caseCount).toBe(0);
+
+      const read = await store.getSuiteById(tenant, suiteA.id);
+      expect(read).toMatchObject({
+        id: suiteA.id,
+        tenantId: tenant,
+        projectId: project.id,
+        name: "Smoke",
+        caseCount: 1,
+      });
+    });
+
+    it("updates a suite and returns null / false for an absent id", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const suite = await createSuite(tenant, project.id, "Before");
+
+      const updated = await store.updateSuite(tenant, suite.id, { name: "After" });
+      expect(updated).toMatchObject({ id: suite.id, name: "After", projectId: project.id });
+
+      const missing = randomUUID();
+      expect(await store.getSuiteById(tenant, missing)).toBeNull();
+      expect(await store.updateSuite(tenant, missing, { name: "nope" })).toBeNull();
+      expect(await store.deleteSuite(tenant, missing)).toBe(false);
+    });
+
+    it("deletes a suite", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const suite = await createSuite(tenant, project.id);
+
+      expect(await store.deleteSuite(tenant, suite.id)).toBe(true);
+      expect(await store.getSuiteById(tenant, suite.id)).toBeNull();
+    });
+  });
+
+  describe("cases", () => {
+    it("round-trips case fields including interruptPolicy", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const suite = await createSuite(tenant, project.id);
+
+      const createdCase = await createCase(tenant, suite.id);
+
+      expect(createdCase).toMatchObject({
+        tenantId: tenant,
+        suiteId: suite.id,
+        inputMessage: "Approve the payment",
+        inputFiles: { receipt: "file-1" },
+        steps: [{ agent_id: "executor-1", override_message: "go" }],
+        outputType: "file_content",
+        contentAssertion: "payment executed",
+        rubrics: [{ name: "accuracy", weight: 0.5, description: "Correct result" }],
+        interruptPolicy: { mode: "canned-response", value: "同意" },
+      });
+
+      const read = await store.getCaseById(tenant, createdCase.id);
+      expect(read).toEqual(createdCase);
+
+      const listed = await store.getCasesBySuite(tenant, suite.id);
+      expect(listed.map((c) => c.id)).toEqual([createdCase.id]);
+    });
+
+    it("applies defaults for omitted optional case fields", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const suite = await createSuite(tenant, project.id);
+
+      const createdCase = await store.createCase(tenant, suite.id, randomUUID(), {
+        inputMessage: "hello",
+        steps: [{ agent_id: "executor-1" }],
+        outputType: "message_content",
+        contentAssertion: "",
+      });
+
+      expect(createdCase.inputFiles).toEqual({});
+      expect(createdCase.rubrics).toEqual([]);
+      expect(createdCase.contentAssertion).toBe("");
+      expect(createdCase.interruptPolicy).toBeUndefined();
+    });
+
+    it("updates and deletes a case", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const suite = await createSuite(tenant, project.id);
+      const createdCase = await createCase(tenant, suite.id);
+
+      const updated = await store.updateCase(tenant, createdCase.id, {
+        inputMessage: "updated",
+        interruptPolicy: { mode: "auto-approve" },
+      });
+
+      expect(updated).toMatchObject({
+        id: createdCase.id,
+        inputMessage: "updated",
+        interruptPolicy: { mode: "auto-approve" },
+        contentAssertion: "payment executed",
+      });
+
+      expect(await store.deleteCase(tenant, createdCase.id)).toBe(true);
+      expect(await store.getCaseById(tenant, createdCase.id)).toBeNull();
+
+      const missing = randomUUID();
+      expect(await store.getCaseById(tenant, missing)).toBeNull();
+      expect(await store.updateCase(tenant, missing, { inputMessage: "nope" })).toBeNull();
+      expect(await store.deleteCase(tenant, missing)).toBe(false);
+    });
+  });
+
+  describe("runs", () => {
+    it("creates runs and filters them by project and status", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const otherProject = await createProject(tenant, "Other");
+      const runA = await createRun(tenant, project.id);
+      const runB = await createRun(tenant, project.id);
+      const runC = await createRun(tenant, otherProject.id);
+      await store.updateRunStatus(tenant, runB.id, { status: "completed" });
+
+      expect(runA.status).toBe("running");
+      expect(runA.totalCases).toBe(4);
+      expect(runA.concurrency).toBe(2);
+      expect(runA.holdout).toBe(true);
+      expect(runA.envProjectId).toBe("env-project");
+      expect(runA.envWorkspaceId).toBe("env-workspace");
+      expect(runA.taskId).toBe("task-1");
+      expect(runA.startedAt).toBeInstanceOf(Date);
+
+      const byProject = await store.getRunsByTenant(tenant, { projectId: project.id });
+      expect(byProject.map((r) => r.id).sort()).toEqual([runA.id, runB.id].sort());
+
+      const completed = await store.getRunsByTenant(tenant, { status: "completed" });
+      expect(completed.map((r) => r.id)).toEqual([runB.id]);
+      expect(completed[0].id).not.toBe(runC.id);
+    });
+
+    it("updates run status counters and timestamps", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const run = await createRun(tenant, project.id);
+      const completedAt = new Date("2026-01-02T03:04:05.000Z");
+
+      const updated = await store.updateRunStatus(tenant, run.id, {
+        status: "completed",
+        passedCases: 2,
+        failedCases: 1,
+        interruptedCases: 1,
+        avgScore: 0.6,
+        completedAt,
+      });
+
+      expect(updated).toMatchObject({
+        id: run.id,
+        status: "completed",
+        totalCases: 4,
+        passedCases: 2,
+        failedCases: 1,
+        interruptedCases: 1,
+        avgScore: 0.6,
+      });
+      expect(updated!.completedAt?.toISOString()).toBe(completedAt.toISOString());
+
+      const read = await store.getRunById(tenant, run.id);
+      expect(read).toMatchObject({
+        status: "completed",
+        passedCases: 2,
+        interruptedCases: 1,
+        avgScore: 0.6,
+      });
+    });
+
+    it("requests cancellation only while the run is running", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const run = await createRun(tenant, project.id);
+
+      expect(await store.requestCancel(tenant, run.id)).toBe(true);
+      expect((await store.getRunById(tenant, run.id))?.cancelRequested).toBe(true);
+
+      await store.updateRunStatus(tenant, run.id, { status: "completed" });
+      expect(await store.requestCancel(tenant, run.id)).toBe(false);
+
+      const failedRun = await createRun(tenant, project.id);
+      await store.updateRunStatus(tenant, failedRun.id, { status: "failed" });
+      expect(await store.requestCancel(tenant, failedRun.id)).toBe(false);
+
+      const abortedRun = await createRun(tenant, project.id);
+      await store.updateRunStatus(tenant, abortedRun.id, { status: "aborted" });
+      expect(await store.requestCancel(tenant, abortedRun.id)).toBe(false);
+
+      expect(await store.requestCancel(tenant, randomUUID())).toBe(false);
+    });
+
+    it("deletes a run together with its results", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const run = await createRun(tenant, project.id);
+      await createResult(tenant, run.id);
+
+      expect(await store.deleteRun(tenant, run.id)).toBe(true);
+      expect(await store.getRunById(tenant, run.id)).toBeNull();
+
+      const remaining = await pool.query<{ count: string }>(
+        "SELECT COUNT(*) AS count FROM lattice_eval_run_results WHERE run_id = $1",
+        [run.id],
+      );
+      expect(remaining.rows[0].count).toBe("0");
+    });
+
+    it("returns null for an absent run id", async () => {
+      const tenant = tenantId();
+      const missing = randomUUID();
+
+      expect(await store.getRunById(tenant, missing)).toBeNull();
+      expect(await store.updateRunStatus(tenant, missing, { status: "completed" })).toBeNull();
+      expect(await store.deleteRun(tenant, missing)).toBe(false);
+    });
+  });
+
+  describe("run results", () => {
+    it("creates and lists run results", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const run = await createRun(tenant, project.id);
+
+      const result = await createResult(tenant, run.id);
+      expect(result).toMatchObject({
+        runId: run.id,
+        suiteName: "main",
+        pass: true,
+        score: 0.9,
+        summary: "looks good",
+        durationMs: 120,
+        interrupted: false,
+      });
+      expect(result.dimensionResults).toEqual([
+        { name: "quality", score: 0.9, reason: "clear" },
+      ]);
+      expect(result.messages).toEqual([{ role: "assistant", content: "done", id: "m1" }]);
+      expect(result.logs).toEqual([
+        { timestamp: "2026-01-01T00:00:00.000Z", level: "info", message: "ran" },
+      ]);
+
+      const listed = await store.getResultsByRun(tenant, run.id);
+      expect(listed.map((r) => r.id)).toEqual([result.id]);
+    });
+
+    it("updates and deletes a run result", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const run = await createRun(tenant, project.id);
+      const result = await createResult(tenant, run.id);
+
+      const updated = await store.updateRunResult(tenant, result.id, {
+        pass: false,
+        score: 0.25,
+        summary: "needs work",
+        durationMs: 500,
+      });
+
+      expect(updated).toMatchObject({
+        id: result.id,
+        runId: run.id,
+        pass: false,
+        score: 0.25,
+        summary: "needs work",
+        durationMs: 500,
+      });
+
+      expect(await store.deleteRunResult(tenant, result.id)).toBe(true);
+      expect(await store.deleteRunResult(tenant, result.id)).toBe(false);
+      expect(await store.getResultsByRun(tenant, run.id)).toEqual([]);
+    });
+
+    it("returns the updated interrupted flag from updateRunResult", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const run = await createRun(tenant, project.id);
+      const result = await createResult(tenant, run.id);
+
+      const updated = await store.updateRunResult(tenant, result.id, { interrupted: true });
+
+      expect(updated?.interrupted).toBe(true);
+      const [read] = await store.getResultsByRun(tenant, run.id);
+      expect(read.interrupted).toBe(true);
+    });
+
+    it("rejects moving a result to another run", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant);
+      const run = await createRun(tenant, project.id);
+      const result = await createResult(tenant, run.id);
+
+      await expect(
+        store.updateRunResult(tenant, result.id, { runId: randomUUID() }),
+      ).rejects.toThrow("runId cannot be updated on an existing result");
+    });
+  });
+
+  describe("project report", () => {
+    it("averages scores and takes the pass rate from the newest run", async () => {
+      const tenant = tenantId();
+      const project = await createProject(tenant, "Reported");
+      const otherProject = await createProject(tenant, "Unrelated");
+      const olderRun = await createRun(tenant, project.id);
+      const newerRun = await createRun(tenant, project.id);
+      const otherRun = await createRun(tenant, otherProject.id);
+
+      await store.updateRunStatus(tenant, olderRun.id, {
+        status: "completed",
+        passedCases: 1,
+        failedCases: 3,
+        avgScore: 0.25,
+      });
+      await store.updateRunStatus(tenant, newerRun.id, {
+        status: "completed",
+        passedCases: 3,
+        failedCases: 1,
+        avgScore: 0.75,
+      });
+      await store.updateRunStatus(tenant, otherRun.id, {
+        status: "completed",
+        passedCases: 4,
+        avgScore: 1,
+      });
+      await pool.query("UPDATE lattice_eval_runs SET created_at = $1 WHERE id = $2", [
+        new Date("2026-01-01T00:00:00.000Z"),
+        olderRun.id,
+      ]);
+      await pool.query("UPDATE lattice_eval_runs SET created_at = $1 WHERE id = $2", [
+        new Date("2026-01-02T00:00:00.000Z"),
+        newerRun.id,
+      ]);
+
+      const report = await store.getProjectReport(tenant, project.id);
+
+      expect(report).toMatchObject({
+        projectId: project.id,
+        projectName: "Reported",
+        totalRuns: 2,
+        latestPassRate: 0.75,
+        avgScore: 0.5,
+      });
+      expect(report!.runs.map((r) => r.id)).toEqual([newerRun.id, olderRun.id]);
+    });
+  });
+
+  describe("tenant isolation", () => {
+    it("does not expose another tenant's project, case, suite, or run", async () => {
+      const tenantA = tenantId();
+      const tenantB = tenantId();
+      const project = await createProject(tenantA);
+      const suite = await createSuite(tenantA, project.id);
+      const testCase = await createCase(tenantA, suite.id);
+      const run = await createRun(tenantA, project.id);
+
+      expect(await store.getProjectById(tenantB, project.id)).toBeNull();
+      expect(await store.getSuiteById(tenantB, suite.id)).toBeNull();
+      expect(await store.getCaseById(tenantB, testCase.id)).toBeNull();
+      expect(await store.getRunById(tenantB, run.id)).toBeNull();
+      expect(await store.getSuitesByProject(tenantB, project.id)).toEqual([]);
+      expect(await store.getCasesBySuite(tenantB, suite.id)).toEqual([]);
+      expect(await store.getRunsByTenant(tenantB, { projectId: project.id })).toEqual([]);
+      expect(await store.updateProject(tenantB, project.id, { name: "hijack" })).toBeNull();
+      expect(await store.deleteProject(tenantB, project.id)).toBe(false);
+    });
+  });
+});