@axiom-lattice/pg-stores 2.0.6 → 2.0.8

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "2.0.6",
+  "version": "2.0.8",
   "description": "PG stores implementation for Axiom Lattice framework",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -25,8 +25,8 @@
     "@langchain/core": "1.1.30",
     "pg": "^8.16.3",
     "uuid": "^9.0.1",
-    "@axiom-lattice/core": "3.0.6",
-    "@axiom-lattice/protocols": "3.0.2"
+    "@axiom-lattice/core": "3.1.0",
+    "@axiom-lattice/protocols": "3.0.3"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",

package/src/__tests__/eval-task-id.test.ts

@@ -0,0 +1,34 @@
+/**
+ * Eval task_id migration tests
+ *
+ * Verifies the add_task_id_to_eval_runs migration adds (and removes)
+ * the task_id column on lattice_eval_runs.
+ */
+
+import { describe, expect, it, jest } from "@jest/globals";
+import { addTaskIdToEvalRuns, evalMigrations } from "../migrations/eval_migrations";
+
+describe("add_task_id_to_eval_runs migration", () => {
+  it("is registered in evalMigrations", () => {
+    expect(evalMigrations.map((m) => m.name)).toContain("add_task_id_to_eval_runs");
+  });
+
+  it("has a version greater than all existing eval migrations", () => {
+    const versions = evalMigrations.map((m) => m.version);
+    expect(addTaskIdToEvalRuns.version).toBe(Math.max(...versions));
+  });
+
+  it("emits ADD COLUMN task_id in up()", async () => {
+    const queries: string[] = [];
+    const client = { query: jest.fn(async (sql: string) => { queries.push(sql); }) };
+    await addTaskIdToEvalRuns.up(client as never);
+    expect(queries.some((q) => q.includes("ADD COLUMN IF NOT EXISTS task_id TEXT"))).toBe(true);
+  });
+
+  it("emits DROP COLUMN task_id in down()", async () => {
+    const queries: string[] = [];
+    const client = { query: jest.fn(async (sql: string) => { queries.push(sql); }) };
+    await addTaskIdToEvalRuns.down(client as never);
+    expect(queries.some((q) => q.includes("DROP COLUMN IF EXISTS task_id"))).toBe(true);
+  });
+});

package/src/__tests__/task-files.test.ts

@@ -0,0 +1,34 @@
+/**
+ * Task files migration tests
+ *
+ * Verifies the add_files_to_tasks migration adds (and removes)
+ * the files column on lattice_tasks.
+ */
+
+import { describe, expect, it, jest } from "@jest/globals";
+import { addFilesToTasks, taskMigrations } from "../migrations/task_migration";
+
+describe("add_files_to_tasks migration", () => {
+  it("is registered in taskMigrations", () => {
+    expect(taskMigrations.map((m) => m.name)).toContain("add_files_to_tasks");
+  });
+
+  it("has a version greater than all existing task migrations", () => {
+    const versions = taskMigrations.map((m) => m.version);
+    expect(addFilesToTasks.version).toBe(Math.max(...versions));
+  });
+
+  it("emits ADD COLUMN files JSONB in up()", async () => {
+    const queries: string[] = [];
+    const client = { query: jest.fn(async (sql: string) => { queries.push(sql); }) };
+    await addFilesToTasks.up(client as never);
+    expect(queries.some((q) => q.includes("ADD COLUMN IF NOT EXISTS files JSONB"))).toBe(true);
+  });
+
+  it("emits DROP COLUMN files in down()", async () => {
+    const queries: string[] = [];
+    const client = { query: jest.fn(async (sql: string) => { queries.push(sql); }) };
+    await addFilesToTasks.down(client as never);
+    expect(queries.some((q) => q.includes("DROP COLUMN IF EXISTS files"))).toBe(true);
+  });
+});

package/src/createPgStoreConfig.ts

@@ -63,7 +63,7 @@ import { addWorkspaceProjectToQueue } from "./migrations/add_workspace_project_t
 import { createWorkflowTrackingTables, addStepThreadId, addWorkflowRunsTenantStatusUpdatedIndex } from "./migrations/workflow_tracking_migrations";
 import { evalMigrations } from "./migrations/eval_migrations";
 import { createA2AApiKeysTable } from "./migrations/a2a_api_key_migration";
-import { createTasksTable, addTaskFieldsMigration, addTaskProjectFieldsMigration } from "./migrations/task_migration";
+import { taskMigrations } from "./migrations/task_migration";
 import { createMenuItemsTable } from "./migrations/menu_items_migration";
 import { addFileContentType } from "./migrations/menu_items_add_file_type";
 import { createSharedResourcesTable } from "./migrations/shared_resources_migration";
@@ -115,14 +115,14 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(createA2AApiKeysTable);           // v130
   mm.register(addWorkspaceProjectToQueue);      // v131
   mm.register(addAssistantOwnerUserId);         // v132
-  mm.register(createTasksTable);                 // v133
+  for (const m of taskMigrations) {             // v133, v139, v140, v141
+    mm.register(m);
+  }
   mm.register(createMenuItemsTable);             // v134
   mm.register(createSharedResourcesTable);       // v135
   mm.register(addFileContentType);               // v136
   mm.register(createCollectionsTable);           // v137
-  mm.register(addTaskFieldsMigration);            // v139
   mm.register(createTaskWorkItemsMigration);      // v138
-  mm.register(addTaskProjectFieldsMigration);     // v140
   mm.register(addWorkItemProjectFieldsMigration); // v140
   mm.register(createConnectionConfigsTable);      // v160
   mm.register(addWorkflowRunsTenantStatusUpdatedIndex); // v161

package/src/migrations/eval_migrations.ts

@@ -359,6 +359,28 @@ export const enforceUniqueEvalProjectName: Migration = {
   },
 };
 
+/** Add task_id to eval runs (training round association, 1 task : N runs) */
+export const addTaskIdToEvalRuns: Migration = {
+  version: 164,
+  name: "add_task_id_to_eval_runs",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_eval_runs
+        ADD COLUMN IF NOT EXISTS task_id TEXT
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_eval_runs_task_id
+        ON lattice_eval_runs(task_id)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query(`DROP INDEX IF EXISTS idx_lattice_eval_runs_task_id`);
+    await client.query(`
+      ALTER TABLE lattice_eval_runs DROP COLUMN IF EXISTS task_id
+    `);
+  },
+};
+
 /** All eval migrations in version order */
 export const evalMigrations: Migration[] = [
   createEvalProjectsTable,
@@ -371,4 +393,5 @@ export const evalMigrations: Migration[] = [
   addInterruptColumnsToEval,
   addInterruptPolicyToEvalCases,
   enforceUniqueEvalProjectName,
+  addTaskIdToEvalRuns,
 ];

package/src/migrations/task_migration.ts

@@ -80,3 +80,28 @@ export const addTaskProjectFieldsMigration: Migration = {
     `);
   },
 };
+
+/** Add files column to tasks (URI file references) */
+export const addFilesToTasks: Migration = {
+  version: 141,
+  name: "add_files_to_tasks",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_tasks
+        ADD COLUMN IF NOT EXISTS files JSONB
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_tasks DROP COLUMN IF EXISTS files
+    `);
+  },
+};
+
+/** All task migrations in version order */
+export const taskMigrations: Migration[] = [
+  createTasksTable,
+  addTaskFieldsMigration,
+  addTaskProjectFieldsMigration,
+  addFilesToTasks,
+];

package/src/stores/PostgreSQLEvalStore.ts

@@ -183,6 +183,7 @@ export class PostgreSQLEvalStore implements EvalStore {
       holdout: row.holdout as boolean | undefined,
       envProjectId: row.env_project_id as string | undefined,
       envWorkspaceId: row.env_workspace_id as string | undefined,
+      taskId: row.task_id as string | undefined,
       createdAt: new Date(row.created_at as string),
       startedAt: row.started_at ? new Date(row.started_at as string) : undefined,
       completedAt: row.completed_at ? new Date(row.completed_at as string) : undefined,
@@ -609,7 +610,7 @@ export class PostgreSQLEvalStore implements EvalStore {
       `SELECT id, project_id, tenant_id,
               status, concurrency, total_cases,
               passed_cases, failed_cases, interrupted_cases, avg_score,
-              error, holdout, env_project_id, env_workspace_id, created_at, started_at, completed_at
+              error, holdout, env_project_id, env_workspace_id, task_id, created_at, started_at, completed_at
        FROM lattice_eval_runs
        WHERE ${conditions.join(" AND ")}
        ORDER BY created_at DESC`,
@@ -631,7 +632,7 @@ export class PostgreSQLEvalStore implements EvalStore {
       `SELECT id, project_id, tenant_id,
               status, concurrency, total_cases,
               passed_cases, failed_cases, interrupted_cases, avg_score,
-              error, holdout, env_project_id, env_workspace_id, created_at, started_at, completed_at
+              error, holdout, env_project_id, env_workspace_id, task_id, created_at, started_at, completed_at
        FROM lattice_eval_runs
        WHERE id = $1 AND tenant_id = $2`,
       [id, tenantId]
@@ -656,12 +657,12 @@ export class PostgreSQLEvalStore implements EvalStore {
       completed_at: string | null;
     }>(
       `INSERT INTO lattice_eval_runs
-        (id, project_id, tenant_id, status, concurrency, total_cases, passed_cases, failed_cases, avg_score, holdout, env_project_id, env_workspace_id, started_at)
-       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
+        (id, project_id, tenant_id, status, concurrency, total_cases, passed_cases, failed_cases, avg_score, holdout, env_project_id, env_workspace_id, task_id, started_at)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
        RETURNING id, project_id, tenant_id,
                  status, concurrency, total_cases,
                  passed_cases, failed_cases, interrupted_cases, avg_score,
-                 error, holdout, env_project_id, env_workspace_id, created_at, started_at, completed_at`,
+                 error, holdout, env_project_id, env_workspace_id, task_id, created_at, started_at, completed_at`,
       [
         actualId,
         projectId,
@@ -675,6 +676,7 @@ export class PostgreSQLEvalStore implements EvalStore {
         data.holdout ?? false,
         data.envProjectId || null,
         data.envWorkspaceId || null,
+        data.taskId || null,
         new Date(),
       ]
     );
@@ -720,7 +722,7 @@ export class PostgreSQLEvalStore implements EvalStore {
        RETURNING id, project_id, tenant_id,
                  status, concurrency, total_cases,
                  passed_cases, failed_cases, interrupted_cases, avg_score,
-                 error, holdout, env_project_id, env_workspace_id, created_at, started_at, completed_at`,
+                 error, holdout, env_project_id, env_workspace_id, task_id, created_at, started_at, completed_at`,
       vals
     );
     return rows.length ? this.mapRowToRun(rows[0] as unknown as Record<string, unknown>) : null;

package/src/stores/PostgreSQLTaskStore.ts

@@ -10,9 +10,10 @@ import type {
   CreateTaskRequest,
   UpdateTaskRequest,
   TaskListFilter,
+  TaskFileRef,
 } from "@axiom-lattice/protocols";
 import { MigrationManager } from "../migrations/migration";
-import { createTasksTable, addTaskFieldsMigration, addTaskProjectFieldsMigration } from "../migrations/task_migration";
+import { createTasksTable, addTaskFieldsMigration, addTaskProjectFieldsMigration, addFilesToTasks } from "../migrations/task_migration";
 import { v4 as uuidv4 } from "uuid";
 
 interface TaskRow {
@@ -35,10 +36,22 @@ interface TaskRow {
   failure_reason: string | null;
   workspace_id: string | null;
   project_id: string | null;
+  files: unknown;
   created_at: string;
   updated_at: string;
 }
 
+function parseTaskFiles(raw: unknown): TaskFileRef[] | undefined {
+  if (!raw) return undefined;
+  try {
+    const parsed = typeof raw === "string" ? JSON.parse(raw) : raw;
+    if (!Array.isArray(parsed)) return undefined;
+    return parsed as TaskFileRef[];
+  } catch {
+    return undefined;
+  }
+}
+
 function mapRowToTask(row: TaskRow): TaskItem {
   return {
     id: row.id,
@@ -60,6 +73,7 @@ function mapRowToTask(row: TaskRow): TaskItem {
     failureReason: row.failure_reason ?? undefined,
     workspaceId: row.workspace_id ?? undefined,
     projectId: row.project_id ?? undefined,
+    files: parseTaskFiles(row.files),
     createdAt: new Date(row.created_at),
     updatedAt: new Date(row.updated_at),
   };
@@ -99,6 +113,7 @@ export class PostgreSQLTaskStore implements TaskStore {
     this.migrationManager.register(createTasksTable);
     this.migrationManager.register(addTaskFieldsMigration);
     this.migrationManager.register(addTaskProjectFieldsMigration);
+    this.migrationManager.register(addFilesToTasks);
 
     if (options.autoMigrate !== false) {
       this.initialize().catch((error) => {
@@ -137,8 +152,8 @@ export class PostgreSQLTaskStore implements TaskStore {
     const now = new Date().toISOString();
 
     await this.pool.query(
-      `INSERT INTO lattice_tasks (id, tenant_id, owner_type, owner_id, title, description, status, priority, due_date, metadata, parent_id, source_id, context, require_review, dependencies, result, failure_reason, workspace_id, project_id, created_at, updated_at)
-       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21)`,
+      `INSERT INTO lattice_tasks (id, tenant_id, owner_type, owner_id, title, description, status, priority, due_date, metadata, parent_id, source_id, context, require_review, dependencies, result, failure_reason, workspace_id, project_id, files, created_at, updated_at)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22)`,
       [
         id,
         params.tenantId,
@@ -159,6 +174,7 @@ export class PostgreSQLTaskStore implements TaskStore {
         params.failureReason ?? null,
         params.workspaceId || null,
         params.projectId || null,
+        params.files ? JSON.stringify(params.files) : null,
         now,
         now,
       ],
@@ -276,6 +292,10 @@ export class PostgreSQLTaskStore implements TaskStore {
       setClauses.push(`metadata = $${paramIndex++}`);
       params.push(JSON.stringify(updates.metadata));
     }
+    if (updates.files !== undefined) {
+      setClauses.push(`files = $${paramIndex++}`);
+      params.push(JSON.stringify(updates.files));
+    }
     if (updates.parentId !== undefined) {
       setClauses.push(`parent_id = $${paramIndex++}`);
       params.push(updates.parentId);