@axiom-lattice/pg-stores 2.0.9 → 3.0.0

package/src/__tests__/PostgreSQLTaskWorkItemStore.test.ts

@@ -0,0 +1,104 @@
+import type { Pool } from "pg";
+import { PostgreSQLTaskWorkItemStore } from "../stores/PostgreSQLTaskWorkItemStore";
+import { addTaskWorkItemEventKeyMigration } from "../migrations/task_work_items_migration";
+
+describe("PostgreSQLTaskWorkItemStore", () => {
+  const row = {
+    id: "item-1", tenant_id: "tenant-1", task_id: "task-1", event_key: "evidence:key",
+    action: "completion_evidence", actor: "agent:a", thread_id: null, summary: "first",
+    detail: null, attempt: null, workspace_id: null, project_id: null,
+    created_at: "2026-01-01T00:00:00.000Z",
+  };
+
+  it.each([
+    [undefined, "ASC"],
+    ["asc", "ASC"],
+    ["desc", "DESC"],
+  ] as const)("uses %s order as %s with deterministic IDs", async (order, sqlOrder) => {
+    const query = jest.fn().mockResolvedValue({ rows: [] });
+    const store = new PostgreSQLTaskWorkItemStore({ query } as unknown as Pool);
+
+    await store.list({
+      tenantId: "tenant-1",
+      taskId: "task-1",
+      limit: 100,
+      ...(order && { order }),
+    });
+
+    expect(query).toHaveBeenCalledWith(
+      expect.stringContaining(`ORDER BY created_at ${sqlOrder}, id ${sqlOrder} LIMIT $3`),
+      ["tenant-1", "task-1", 100],
+    );
+  });
+
+  it("filters by workspaceId and projectId", async () => {
+    const query = jest.fn().mockResolvedValue({ rows: [] });
+    const store = new PostgreSQLTaskWorkItemStore({ query } as unknown as Pool);
+
+    await store.list({
+      tenantId: "tenant-1",
+      taskId: "task-1",
+      workspaceId: "workspace-1",
+      projectId: "project-1",
+      order: "desc",
+      limit: 100,
+    });
+
+    expect(query).toHaveBeenCalledWith(
+      expect.stringContaining(
+        "workspace_id = $3 AND project_id = $4 ORDER BY created_at DESC, id DESC LIMIT $5",
+      ),
+      ["tenant-1", "task-1", "workspace-1", "project-1", 100],
+    );
+  });
+
+  it("looks up an event key by tenant and task", async () => {
+    const query = jest.fn().mockResolvedValue({ rows: [row] });
+    const store = new PostgreSQLTaskWorkItemStore({ query } as unknown as Pool);
+
+    const result = await store.findByEventKey("tenant-1", "task-1", "evidence:key");
+
+    expect(query).toHaveBeenCalledWith(
+      expect.stringContaining("tenant_id = $1 AND task_id = $2 AND event_key = $3"),
+      ["tenant-1", "task-1", "evidence:key"],
+    );
+    expect(result?.eventKey).toBe("evidence:key");
+  });
+
+  it("maps a null event key to undefined", async () => {
+    const query = jest.fn().mockResolvedValue({ rows: [{ ...row, event_key: null }] });
+    const store = new PostgreSQLTaskWorkItemStore({ query } as unknown as Pool);
+
+    const result = await store.list({ tenantId: "tenant-1", taskId: "task-1" });
+
+    expect(result[0].eventKey).toBeUndefined();
+  });
+
+  it("atomically inserts by task-scoped event key or returns the existing row", async () => {
+    const query = jest.fn().mockResolvedValue({ rows: [row] });
+    const store = new PostgreSQLTaskWorkItemStore({ query } as unknown as Pool);
+
+    const result = await store.createIfAbsentByEventKey({
+      tenantId: "tenant-1", taskId: "task-1", eventKey: "evidence:key",
+      action: "completion_evidence", actor: "agent:a", summary: "first",
+    });
+
+    expect(query).toHaveBeenCalledTimes(1);
+    expect(query).toHaveBeenCalledWith(
+      expect.stringMatching(/INSERT[\s\S]+ON CONFLICT \(tenant_id, task_id, event_key\)[\s\S]+RETURNING/),
+      expect.arrayContaining(["tenant-1", "task-1", "evidence:key"]),
+    );
+    expect(result.eventKey).toBe("evidence:key");
+  });
+
+  it("migrates a task-scoped unique event key", async () => {
+    const query = jest.fn().mockResolvedValue({ rows: [] });
+
+    await addTaskWorkItemEventKeyMigration.up({ query } as never);
+
+    expect(query).toHaveBeenCalledWith(expect.stringContaining("ADD COLUMN IF NOT EXISTS event_key TEXT"));
+    expect(query).toHaveBeenCalledWith(expect.stringContaining(
+      "(tenant_id, task_id, event_key) WHERE event_key IS NOT NULL",
+    ));
+  });
+});

package/src/__tests__/add_a2a_key_assistant_ids.test.ts

@@ -0,0 +1,34 @@
+/**
+ * add_a2a_key_assistant_ids migration tests
+ *
+ * Verifies the migration adds (and removes) the assistant_ids JSONB
+ * column on lattice_a2a_api_keys.
+ */
+
+import { describe, expect, it, jest } from "@jest/globals";
+import type { PoolClient } from "pg";
+import { addA2AKeyAssistantIds } from "../migrations/add_a2a_key_assistant_ids";
+
+describe("add_a2a_key_assistant_ids migration", () => {
+  it("has name add_a2a_key_assistant_ids", () => {
+    expect(addA2AKeyAssistantIds.name).toBe("add_a2a_key_assistant_ids");
+  });
+
+  it("has a version greater than the current max (166)", () => {
+    expect(addA2AKeyAssistantIds.version).toBe(167);
+  });
+
+  it("emits ADD COLUMN assistant_ids JSONB in up()", async () => {
+    const queries: string[] = [];
+    const client = { query: jest.fn(async (sql: string) => { queries.push(sql); }) };
+    await addA2AKeyAssistantIds.up(client as unknown as PoolClient);
+    expect(queries.some((q) => q.includes("ADD COLUMN IF NOT EXISTS assistant_ids JSONB"))).toBe(true);
+  });
+
+  it("emits DROP COLUMN assistant_ids in down()", async () => {
+    const queries: string[] = [];
+    const client = { query: jest.fn(async (sql: string) => { queries.push(sql); }) };
+    await addA2AKeyAssistantIds.down?.(client as unknown as PoolClient);
+    expect(queries.some((q) => q.includes("DROP COLUMN IF EXISTS assistant_ids"))).toBe(true);
+  });
+});

package/src/__tests__/workspace-project-store.test.ts

@@ -232,4 +232,52 @@ describeIfDb("PostgreSQLProjectStore", () => {
     const deleted = await store.deleteProject(tenantId, "non-existent");
     expect(deleted).toBe(false);
   });
+
+  test("should default kind to 'business' when omitted on create", async () => {
+    const project = await store.createProject(tenantId, workspaceId, "test-proj-kind-1", {
+      name: "Default Kind",
+    });
+    expect(project.kind).toBe("business");
+
+    const fetched = await store.getProjectById(tenantId, "test-proj-kind-1");
+    expect(fetched?.kind).toBe("business");
+  });
+
+  test("should round-trip an explicit kind through create and get", async () => {
+    const project = await store.createProject(tenantId, workspaceId, "test-proj-kind-2", {
+      name: "Training Project",
+      kind: "training",
+    });
+    expect(project.kind).toBe("training");
+
+    const fetched = await store.getProjectById(tenantId, "test-proj-kind-2");
+    expect(fetched?.kind).toBe("training");
+  });
+
+  test("should update kind", async () => {
+    await store.createProject(tenantId, workspaceId, "test-proj-kind-3", {
+      name: "Reclassified",
+    });
+
+    const updated = await store.updateProject(tenantId, "test-proj-kind-3", { kind: "personal" });
+    expect(updated?.kind).toBe("personal");
+  });
+
+  test("should filter projects by kind", async () => {
+    await store.createProject(tenantId, workspaceId, "test-proj-kind-4", { name: "Biz", kind: "business" });
+    await store.createProject(tenantId, workspaceId, "test-proj-kind-5", { name: "Train", kind: "training" });
+
+    const training = await store.getProjectsByWorkspace(tenantId, workspaceId, { kind: "training" });
+    const ids = training.filter((p) => p.id.startsWith("test-proj-kind-")).map((p) => p.id);
+    expect(ids).toEqual(["test-proj-kind-5"]);
+  });
+
+  test("should return all projects when kind filter is omitted", async () => {
+    await store.createProject(tenantId, workspaceId, "test-proj-kind-6", { name: "Biz", kind: "business" });
+    await store.createProject(tenantId, workspaceId, "test-proj-kind-7", { name: "Train", kind: "training" });
+
+    const all = await store.getProjectsByWorkspace(tenantId, workspaceId);
+    const ids = all.filter((p) => p.id.startsWith("test-proj-kind-")).map((p) => p.id);
+    expect(ids).toHaveLength(2);
+  });
 });

package/src/createPgStoreConfig.ts

@@ -43,6 +43,7 @@ import { changeSkillPrimaryKey } from "./migrations/skill_pk_migration";
 import { createWorkspacesTable } from "./migrations/workspace_migrations";
 import { createProjectsTable } from "./migrations/project_migrations";
 import { addProjectConfigColumn } from "./migrations/add_project_config_column";
+import { addProjectKindColumn } from "./migrations/add_project_kind_column";
 import { createUsersTable } from "./migrations/user_migrations";
 import { addUserStatusColumn } from "./migrations/user_status_migration";
 import { createTenantsTable } from "./migrations/tenant_migrations";
@@ -63,12 +64,13 @@ 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 { addA2AKeyAssistantIds } from "./migrations/add_a2a_key_assistant_ids";
 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";
 import { createCollectionsTable } from "./migrations/collection_migrations";
-import { createTaskWorkItemsMigration, addWorkItemProjectFieldsMigration } from "./migrations/task_work_items_migration";
+import { createTaskWorkItemsMigration, addWorkItemProjectFieldsMigration, addTaskWorkItemEventKeyMigration } from "./migrations/task_work_items_migration";
 
 export async function createPgStoreConfig(connectionString: string) {
   const pool = new Pool({ connectionString });
@@ -126,6 +128,9 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(addWorkItemProjectFieldsMigration); // v140
   mm.register(createConnectionConfigsTable);      // v160
   mm.register(addWorkflowRunsTenantStatusUpdatedIndex); // v161
+  mm.register(addProjectKindColumn);          // v165
+  mm.register(addA2AKeyAssistantIds);         // v167
+  mm.register(addTaskWorkItemEventKeyMigration); // v168
 
   await mm.migrate();
 

package/src/migrations/add_a2a_key_assistant_ids.ts

@@ -0,0 +1,24 @@
+import type { Migration } from "./migration";
+import type { PoolClient } from "pg";
+
+/**
+ * Adds the assistant_ids JSONB whitelist column to A2A API keys.
+ *
+ * Empty/undefined assistant_ids means all exposed agents in the tenant.
+ */
+export const addA2AKeyAssistantIds: Migration = {
+  version: 167,
+  name: "add_a2a_key_assistant_ids",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_a2a_api_keys
+        ADD COLUMN IF NOT EXISTS assistant_ids JSONB
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_a2a_api_keys
+        DROP COLUMN IF EXISTS assistant_ids
+    `);
+  },
+};

package/src/migrations/add_project_kind_column.ts

@@ -0,0 +1,29 @@
+/**
+ * Migration: Add kind column to lattice_projects
+ *
+ * Adds a TEXT `kind` column (NOT NULL DEFAULT 'business') to the
+ * `lattice_projects` table for project classification
+ * ('business' | 'training' | 'personal'). Existing rows are backfilled
+ * with 'business' via the column default.
+ */
+
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+/** Adds kind TEXT column to lattice_projects */
+export const addProjectKindColumn: Migration = {
+  version: 165,
+  name: "add_project_kind_column",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_projects
+      ADD COLUMN IF NOT EXISTS kind TEXT NOT NULL DEFAULT 'business'
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_projects
+      DROP COLUMN IF EXISTS kind
+    `);
+  },
+};

package/src/migrations/task_work_items_migration.ts

@@ -49,3 +49,18 @@ export const addWorkItemProjectFieldsMigration: Migration = {
     `);
   },
 };
+
+export const addTaskWorkItemEventKeyMigration: Migration = {
+  version: 168,
+  name: "add_task_work_item_event_key",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_task_work_items
+        ADD COLUMN IF NOT EXISTS event_key TEXT
+    `);
+    await client.query(`
+      CREATE UNIQUE INDEX IF NOT EXISTS idx_task_work_items_event_key
+        ON lattice_task_work_items (tenant_id, task_id, event_key) WHERE event_key IS NOT NULL
+    `);
+  },
+};

package/src/stores/PostgreSQLA2AApiKeyStore.ts

@@ -8,6 +8,7 @@ import type {
 } from "@axiom-lattice/protocols";
 import { MigrationManager } from "../migrations/migration";
 import { createA2AApiKeysTable } from "../migrations/a2a_api_key_migration";
+import { addA2AKeyAssistantIds } from "../migrations/add_a2a_key_assistant_ids";
 import { encrypt, decrypt } from "@axiom-lattice/core";
 import { randomUUID } from "crypto";
 
@@ -21,8 +22,8 @@ type KeyRow = {
   id: string;
   key_value: string;
   tenant_id: string;
-  project_id: string | null;
-  workspace_id: string | null;
+  project_id: string;
+  assistant_ids: string[] | null;
   label: string | null;
   enabled: boolean;
   created_at: Date;
@@ -38,8 +39,8 @@ function mapRowToRecord(row: KeyRow): A2AApiKeyRecord {
     id: row.id,
     key: decrypt(row.key_value),
     tenantId: row.tenant_id,
-    projectId: row.project_id || undefined,
-    workspaceId: row.workspace_id || undefined,
+    projectId: row.project_id,
+    assistantIds: row.assistant_ids ?? undefined,
     label: row.label || undefined,
     enabled: row.enabled,
     createdAt: row.created_at,
@@ -71,7 +72,12 @@ export class PostgreSQLA2AApiKeyStore implements A2AApiKeyStore {
           : (() => { throw new Error("Either pool or poolConfig must be provided"); })();
 
     this.migrationManager = new MigrationManager(this.pool);
+    // Standalone deployments (own poolConfig) must apply every migration the
+    // queries below depend on. `createPgStoreConfig` registers the same
+    // migrations on its own MigrationManager; managers dedupe by name, so the
+    // overlap is safe.
     this.migrationManager.register(createA2AApiKeysTable);
+    this.migrationManager.register(addA2AKeyAssistantIds);
 
     if (options.autoMigrate !== false) {
       this.initialize().catch((error) => {
@@ -116,6 +122,15 @@ export class PostgreSQLA2AApiKeyStore implements A2AApiKeyStore {
     return null;
   }
 
+  async findById(id: string): Promise<A2AApiKeyRecord | null> {
+    await this.ensureInitialized();
+    const result = await this.pool.query<KeyRow>(
+      `SELECT * FROM lattice_a2a_api_keys WHERE id = $1`,
+      [id],
+    );
+    return result.rows[0] ? mapRowToRecord(result.rows[0]) : null;
+  }
+
   async list(params: { tenantId?: string; limit?: number; offset?: number }): Promise<A2AApiKeyRecord[]> {
     await this.ensureInitialized();
     const limit = params.limit || 100;
@@ -138,9 +153,9 @@ export class PostgreSQLA2AApiKeyStore implements A2AApiKeyStore {
     await this.ensureInitialized();
     const key = generateApiKey();
     const result = await this.pool.query<KeyRow>(
-      `INSERT INTO lattice_a2a_api_keys (key_value, tenant_id, project_id, workspace_id, label)
+      `INSERT INTO lattice_a2a_api_keys (key_value, tenant_id, project_id, assistant_ids, label)
        VALUES ($1, $2, $3, $4, $5) RETURNING *`,
-      [encrypt(key), input.tenantId, input.projectId || null, input.workspaceId || null, input.label || null],
+      [encrypt(key), input.tenantId, input.projectId, input.assistantIds ? JSON.stringify(input.assistantIds) : null, input.label || null],
     );
     const record = mapRowToRecord(result.rows[0]);
     record.key = key;
@@ -196,8 +211,8 @@ export class PostgreSQLA2AApiKeyStore implements A2AApiKeyStore {
       map.set(key, {
         key,
         tenantId: row.tenant_id,
-        projectId: row.project_id || undefined,
-        workspaceId: row.workspace_id || undefined,
+        projectId: row.project_id,
+        assistantIds: row.assistant_ids ?? undefined,
       });
     }
     return map;

package/src/stores/PostgreSQLProjectStore.ts

@@ -7,6 +7,7 @@ import type { PoolConfig } from "pg";
 import {
   ProjectStore,
   Project,
+  ProjectFilter,
   CreateProjectRequest,
   UpdateProjectRequest,
 } from "@axiom-lattice/protocols";
@@ -14,6 +15,7 @@ import { MigrationManager } from "../migrations/migration";
 import { createProjectsTable } from "../migrations/project_migrations";
 import { addProjectConfigColumn } from "../migrations/add_project_config_column";
 import { dropProjectFolderPathColumn } from "../migrations/drop_project_folder_path";
+import { addProjectKindColumn } from "../migrations/add_project_kind_column";
 
 /**
  * PostgreSQL ProjectStore options
@@ -69,6 +71,7 @@ export class PostgreSQLProjectStore implements ProjectStore {
     this.migrationManager.register(createProjectsTable);
     this.migrationManager.register(addProjectConfigColumn);
     this.migrationManager.register(dropProjectFolderPathColumn);
+    this.migrationManager.register(addProjectKindColumn);
 
     // Auto-migrate by default
     if (options.autoMigrate !== false) {
@@ -120,6 +123,7 @@ export class PostgreSQLProjectStore implements ProjectStore {
     name: string;
     description: string | null;
     config: unknown | null;
+    kind: string | null;
     created_at: Date;
     updated_at: Date;
   }): Project {
@@ -130,6 +134,7 @@ export class PostgreSQLProjectStore implements ProjectStore {
       name: row.name,
       description: row.description || undefined,
       config: row.config as Record<string, unknown> || undefined,
+      kind: (row.kind || "business") as Project["kind"],
       createdAt: row.created_at,
       updatedAt: row.updated_at,
     };
@@ -138,9 +143,17 @@ export class PostgreSQLProjectStore implements ProjectStore {
   /**
    * Get all projects for a specific workspace
    */
-  async getProjectsByWorkspace(tenantId: string, workspaceId: string): Promise<Project[]> {
+  async getProjectsByWorkspace(tenantId: string, workspaceId: string, filter?: ProjectFilter): Promise<Project[]> {
     await this.ensureInitialized();
 
+    const conditions = ["tenant_id = $1", "workspace_id = $2"];
+    const values: unknown[] = [tenantId, workspaceId];
+
+    if (filter?.kind !== undefined) {
+      values.push(filter.kind);
+      conditions.push(`kind = $${values.length}`);
+    }
+
     const result = await this.pool.query<{
       id: string;
       tenant_id: string;
@@ -148,16 +161,17 @@ export class PostgreSQLProjectStore implements ProjectStore {
       name: string;
       description: string | null;
       config: unknown | null;
+      kind: string | null;
       created_at: Date;
       updated_at: Date;
     }>(
       `
-      SELECT id, tenant_id, workspace_id, name, description, config, created_at, updated_at
+      SELECT id, tenant_id, workspace_id, name, description, config, kind, created_at, updated_at
       FROM lattice_projects
-      WHERE tenant_id = $1 AND workspace_id = $2
+      WHERE ${conditions.join(" AND ")}
       ORDER BY created_at DESC
     `,
-      [tenantId, workspaceId]
+      values
     );
 
     return result.rows.map(this.mapRowToProject);
@@ -176,11 +190,12 @@ export class PostgreSQLProjectStore implements ProjectStore {
       name: string;
       description: string | null;
       config: unknown | null;
+      kind: string | null;
       created_at: Date;
       updated_at: Date;
     }>(
       `
-      SELECT id, tenant_id, workspace_id, name, description, config, created_at, updated_at
+      SELECT id, tenant_id, workspace_id, name, description, config, kind, created_at, updated_at
       FROM lattice_projects
       WHERE id = $1 AND tenant_id = $2
     `,
@@ -206,19 +221,21 @@ export class PostgreSQLProjectStore implements ProjectStore {
     await this.ensureInitialized();
 
     const now = new Date();
+    const kind = data.kind || "business";
 
     await this.pool.query(
       `
-      INSERT INTO lattice_projects (id, tenant_id, workspace_id, name, description, config, created_at, updated_at)
-      VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+      INSERT INTO lattice_projects (id, tenant_id, workspace_id, name, description, config, kind, created_at, updated_at)
+      VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
       ON CONFLICT (id, tenant_id) DO UPDATE SET
         workspace_id = EXCLUDED.workspace_id,
         name = EXCLUDED.name,
         description = EXCLUDED.description,
         config = EXCLUDED.config,
+        kind = EXCLUDED.kind,
         updated_at = EXCLUDED.updated_at
     `,
-      [id, tenantId, workspaceId, data.name, data.description || null, data.config || null, now, now]
+      [id, tenantId, workspaceId, data.name, data.description || null, data.config || null, kind, now, now]
     );
 
     return {
@@ -228,6 +245,7 @@ export class PostgreSQLProjectStore implements ProjectStore {
       name: data.name,
       description: data.description,
       config: data.config,
+      kind,
       createdAt: now,
       updatedAt: now,
     };
@@ -273,6 +291,11 @@ export class PostgreSQLProjectStore implements ProjectStore {
       updateValues.push(updates.config || null);
     }
 
+    if (updates.kind !== undefined) {
+      updateFields.push(`kind = $${paramIndex++}`);
+      updateValues.push(updates.kind);
+    }
+
     if (updateFields.length === 0) {
       // No fields to update
       return existing;