@axiom-lattice/pg-stores 3.0.2 → 3.1.1

package/src/__tests__/ThreadMessageQueueStore.test.ts

@@ -34,6 +34,33 @@ describe('ThreadMessageQueueStore', () => {
   });
 
   describe('addMessage', () => {
+    it('atomically rejects enqueue after the scoped pending capacity is reached', async () => {
+      const client = {
+        query: jest.fn()
+          .mockResolvedValueOnce({})
+          .mockResolvedValueOnce({})
+          .mockResolvedValueOnce({ rows: [{ count: '2' }] })
+          .mockResolvedValueOnce({}),
+        release: jest.fn(),
+      };
+      mockConnect.mockResolvedValueOnce(client);
+
+      await expect(store.addMessageIfCapacity({
+        threadId: mockThreadId,
+        tenantId: mockTenantId,
+        assistantId: mockAssistantId,
+        workspaceId: null,
+        projectId: null,
+        content: 'full',
+      }, 2)).resolves.toBe(false);
+
+      expect(client.query).toHaveBeenCalledWith(
+        'SELECT pg_advisory_xact_lock(hashtext($1))',
+        [`${mockTenantId}:${mockAssistantId}:${mockThreadId}::`],
+      );
+      expect(client.release).toHaveBeenCalled();
+    });
+
     it('should add message to queue', async () => {
       const mockMessage = {
         id: 'msg-1',
@@ -131,7 +158,6 @@ describe('ThreadMessageQueueStore', () => {
     it('should insert message with high priority', async () => {
       // Mock thread info query
       mockQuery
-        .mockResolvedValueOnce({ rows: [{ tenant_id: mockTenantId, assistant_id: mockAssistantId }] }) // Get thread info
         .mockResolvedValueOnce({ rows: [{ next_seq: 1 }] }) // Get next sequence
         .mockResolvedValueOnce({ 
           rows: [{ 
@@ -158,7 +184,6 @@ describe('ThreadMessageQueueStore', () => {
 
     it('should use default tenant/assistant when no existing messages', async () => {
       mockQuery
-        .mockResolvedValueOnce({ rows: [] }) // No existing messages
         .mockResolvedValueOnce({ rows: [{ next_seq: 0 }] }) // Get next sequence
         .mockResolvedValueOnce({ 
           rows: [{ 
@@ -259,6 +284,53 @@ describe('ThreadMessageQueueStore', () => {
       expect(result[0].custom_run_config).toEqual(customRunConfig);
       expect((result[0].content as { command?: unknown }).command).toBeUndefined();
     });
+
+    it('matches nullish dimensions as NULL when a scope is supplied', async () => {
+      mockQuery.mockResolvedValueOnce({ rows: [] });
+
+      await store.getPendingMessages(mockThreadId, {
+        tenantId: mockTenantId,
+        assistantId: mockAssistantId,
+        workspaceId: undefined,
+        projectId: undefined,
+      });
+
+      expect(mockQuery).toHaveBeenCalledWith(
+        expect.stringMatching(/tenant_id = \$2.*assistant_id = \$3.*workspace_id IS NULL.*project_id IS NULL/s),
+        [mockThreadId, mockTenantId, mockAssistantId],
+      );
+    });
+
+    it('applies null-sensitive full scope predicates to every scoped operation', async () => {
+      const scope = {
+        tenantId: mockTenantId,
+        assistantId: mockAssistantId,
+        workspaceId: undefined,
+        projectId: undefined,
+      };
+      mockQuery
+        .mockResolvedValueOnce({ rows: [] })
+        .mockResolvedValueOnce({ rows: [] })
+        .mockResolvedValueOnce({ rows: [{ count: '0' }] })
+        .mockResolvedValueOnce({ rowCount: 0 })
+        .mockResolvedValueOnce({ rowCount: 0 })
+        .mockResolvedValueOnce({ rowCount: 0 })
+        .mockResolvedValueOnce({ rowCount: 0 });
+
+      await store.getProcessingMessages(mockThreadId, scope);
+      await store.getPendingMessages(mockThreadId, scope);
+      await store.getQueueSize(mockThreadId, scope);
+      await store.markProcessing('msg-1', undefined, scope);
+      await store.resetProcessingToPending(mockThreadId, scope);
+      await store.removeMessage('msg-1', scope);
+      await store.clearMessages(mockThreadId, scope);
+
+      expect(mockQuery).toHaveBeenCalledTimes(7);
+      for (const [sql, params] of mockQuery.mock.calls as Array<[string, unknown[]]>) {
+        expect(sql).toMatch(/tenant_id = \$2.*assistant_id = \$3.*workspace_id IS NULL.*project_id IS NULL/s);
+        expect(params).toEqual(expect.arrayContaining([mockTenantId, mockAssistantId]));
+      }
+    });
   });
 
   describe('getQueueSize', () => {
@@ -296,6 +368,26 @@ describe('ThreadMessageQueueStore', () => {
       expect(result[1].threadId).toBe('thread-2');
       expect(result[2].threadId).toBe('thread-3');
     });
+
+    it('preserves NULL workspace and project identity for legacy rows', async () => {
+      mockQuery.mockResolvedValueOnce({
+        rows: [{
+          tenant_id: mockTenantId,
+          assistant_id: mockAssistantId,
+          thread_id: mockThreadId,
+          workspace_id: null,
+          project_id: null,
+        }],
+      });
+
+      await expect(store.getThreadsWithPendingMessages()).resolves.toEqual([{
+        tenantId: mockTenantId,
+        assistantId: mockAssistantId,
+        threadId: mockThreadId,
+        workspaceId: null,
+        projectId: null,
+      }]);
+    });
   });
 
   describe('removeMessage', () => {
@@ -323,7 +415,7 @@ describe('ThreadMessageQueueStore', () => {
       await store.markProcessing('msg-1');
 
       expect(mockQuery).toHaveBeenCalledWith(
-        expect.stringContaining("UPDATE thread_message_queue"),
+        expect.stringContaining("UPDATE lattice_thread_message_queue"),
         ['msg-1']
       );
     });
@@ -336,7 +428,7 @@ describe('ThreadMessageQueueStore', () => {
       await store.clearMessages(mockThreadId);
 
       expect(mockQuery).toHaveBeenCalledWith(
-        expect.stringContaining("DELETE FROM thread_message_queue"),
+        expect.stringContaining("DELETE FROM lattice_thread_message_queue"),
         [mockThreadId]
       );
     });

package/src/__tests__/add_workspace_project_to_queue.test.ts

@@ -0,0 +1,15 @@
+import type { PoolClient } from "pg";
+import { addWorkspaceProjectToQueue } from "../migrations/add_workspace_project_to_queue";
+
+describe("addWorkspaceProjectToQueue", () => {
+  it("adds nullable scope columns so legacy rows retain NULL identity", async () => {
+    const query = jest.fn().mockResolvedValue({ rows: [] });
+
+    await addWorkspaceProjectToQueue.up({ query } as unknown as PoolClient);
+
+    const sql = query.mock.calls.map(([statement]) => statement as string).join("\n");
+    expect(sql).toContain("ADD COLUMN IF NOT EXISTS workspace_id VARCHAR(255)");
+    expect(sql).toContain("ADD COLUMN IF NOT EXISTS project_id VARCHAR(255)");
+    expect(sql).not.toMatch(/SET\s+DEFAULT|NOT\s+NULL|UPDATE\s+lattice_thread_message_queue/i);
+  });
+});

package/src/createPgStoreConfig.ts

@@ -71,8 +71,15 @@ 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, addTaskWorkItemEventKeyMigration } from "./migrations/task_work_items_migration";
+import {
+  createTaskWorkItemsMigration,
+  addWorkItemProjectFieldsMigration,
+  addTaskWorkItemEventKeyMigration,
+  addTaskWorkItemPendingIndexesMigration,
+} from "./migrations/task_work_items_migration";
 import { createAgentWebAppsTable } from "./migrations/agent_web_apps_migration";
+import { createCapabilityBundlesTable } from "./migrations/capability_bundle_migration";
+import { PostgreSQLCapabilityBundleStore } from "./stores/PostgreSQLCapabilityBundleStore";
 
 export async function createPgStoreConfig(connectionString: string) {
   const pool = new Pool({ connectionString });
@@ -134,6 +141,8 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(addA2AKeyAssistantIds);         // v167
   mm.register(addTaskWorkItemEventKeyMigration); // v168
   mm.register(createAgentWebAppsTable);        // v169
+  mm.register(createCapabilityBundlesTable);   // v170
+  mm.register(addTaskWorkItemPendingIndexesMigration); // v171
 
   await mm.migrate();
 
@@ -167,6 +176,7 @@ export async function createPgStoreConfig(connectionString: string) {
     taskWorkItem: taskWorkItemStore,
     a2aApiKey: new PostgreSQLA2AApiKeyStore(opts),
     agentWebApp: new PostgreSQLAgentWebAppStore(opts),
+    capabilityBundle: new PostgreSQLCapabilityBundleStore(opts),
     schedule: new PostgreSQLScheduleStorage(opts),
     menu: new MenuStore(opts),
     sharedResource: new PostgresSharedResourceStore(opts),

package/src/index.ts

@@ -58,6 +58,8 @@ export * from "./stores/ChannelIdentityMappingStore";
 export * from "./stores/PostgreSQLChannelInstallationStore";
 export * from "./stores/PostgreSQLA2AApiKeyStore";
 export * from "./stores/PostgreSQLAgentWebAppStore";
+export * from "./stores/PostgreSQLCapabilityBundleStore";
+export * from "./migrations/capability_bundle_migration";
 export * from "./stores/PostgreSQLWorkflowTrackingStore";
 export * from "./stores/PostgreSQLEvalStore";
 export * from "./migrations/eval_migrations";
@@ -191,4 +193,8 @@ export type {
   AgentWebAppAppearance,
   CreateAgentWebAppInput,
   UpdateAgentWebAppInput,
+  CapabilityBundleStore,
+  CapabilityBundle,
+  CreateCapabilityBundleInput,
+  UpdateCapabilityBundleInput,
 } from "@axiom-lattice/protocols";

package/src/migrations/capability_bundle_migration.ts

@@ -0,0 +1,20 @@
+import type { Migration } from "./migration";
+
+/** Creates the tenant-scoped capability bundle table. */
+export const createCapabilityBundlesTable: Migration = {
+  version: 170,
+  name: "create_capability_bundles_table",
+  up: async (client) => {
+    await client.query(`CREATE TABLE IF NOT EXISTS lattice_capability_bundles (
+      id UUID PRIMARY KEY, tenant_id VARCHAR(255) NOT NULL, bundle_key VARCHAR(255) NOT NULL,
+      name VARCHAR(255) NOT NULL, description TEXT, capabilities JSONB NOT NULL DEFAULT '[]'::jsonb,
+      created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+      UNIQUE (tenant_id, bundle_key)
+    )`);
+    await client.query("CREATE INDEX IF NOT EXISTS idx_lattice_capability_bundles_tenant ON lattice_capability_bundles(tenant_id)");
+  },
+  down: async (client) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_capability_bundles_tenant");
+    await client.query("DROP TABLE IF EXISTS lattice_capability_bundles");
+  },
+};

package/src/migrations/task_work_items_migration.ts

@@ -64,3 +64,19 @@ export const addTaskWorkItemEventKeyMigration: Migration = {
     `);
   },
 };
+
+export const addTaskWorkItemPendingIndexesMigration: Migration = {
+  version: 171,
+  name: "add_task_work_item_pending_indexes",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_task_work_items_pending_order
+        ON lattice_task_work_items (tenant_id, task_id, action, created_at DESC, id DESC)
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_task_work_items_reconciled_result
+        ON lattice_task_work_items (tenant_id, task_id, (detail ->> 'executionResultId'))
+        WHERE action = 'execution_reconciled'
+    `);
+  },
+};

package/src/stores/PostgreSQLCapabilityBundleStore.ts

@@ -0,0 +1,306 @@
+import { randomUUID } from "crypto";
+import { Pool } from "pg";
+import type { PoolConfig } from "pg";
+import type {
+  CapabilityBundle,
+  CapabilityBundleDeleteResult,
+  CapabilityBundleStore,
+  InternalCreateCapabilityBundleInput,
+  InternalUpdateCapabilityBundleInput,
+} from "@axiom-lattice/protocols";
+import { MigrationManager } from "../migrations/migration";
+import { createCapabilityBundlesTable } from "../migrations/capability_bundle_migration";
+
+/**
+ * Configures PostgreSQL connectivity and migration behavior for capability bundles.
+ *
+ * @remarks
+ * An injected `pool` is externally managed and assumed to have migrations applied.
+ * When `poolConfig` is used, the store owns the pool and runs migrations eagerly
+ * unless `autoMigrate` is false. Call `initialize()` before use in that case.
+ */
+export interface PostgreSQLCapabilityBundleStoreOptions {
+  /** Externally managed PostgreSQL pool. */
+  pool?: Pool;
+  /** Connection string or pool configuration used to create an owned pool. */
+  poolConfig?: string | PoolConfig;
+  /** Whether an owned pool should run migrations during construction. @default true */
+  autoMigrate?: boolean;
+}
+
+type Row = {
+  id: string;
+  tenant_id: string;
+  bundle_key: string;
+  name: string;
+  description: string | null;
+  capabilities: unknown;
+  created_at: Date;
+  updated_at: Date | string;
+};
+const duplicateMessage = "Capability bundle key already exists for tenant";
+function map(row: Row): CapabilityBundle {
+  return {
+    id: row.id,
+    tenantId: row.tenant_id,
+    key: row.bundle_key,
+    name: row.name,
+    description: row.description ?? undefined,
+    capabilities: row.capabilities as CapabilityBundle["capabilities"],
+    createdAt: row.created_at.toISOString(),
+    updatedAt: typeof row.updated_at === "string" ? row.updated_at : row.updated_at.toISOString(),
+  };
+}
+
+function isDuplicate(error: unknown): boolean {
+  return typeof error === "object"
+    && error !== null
+    && "code" in error
+    && error.code === "23505";
+}
+
+/**
+ * Persists tenant-scoped capability bundles in PostgreSQL.
+ *
+ * @example
+ * ```ts
+ * const store = new PostgreSQLCapabilityBundleStore({ poolConfig: connectionString });
+ * await store.initialize();
+ * const bundles = await store.listByTenant("tenant-1");
+ * await store.dispose();
+ * ```
+ *
+ * @remarks
+ * Stores created from `poolConfig` own their pool and must be disposed. Stores
+ * created with an injected pool neither migrate nor close that pool.
+ */
+export class PostgreSQLCapabilityBundleStore implements CapabilityBundleStore {
+  private readonly pool: Pool;
+  private migrationManager!: MigrationManager;
+  private initialized = false;
+  private ownsPool = true;
+  private initPromise: Promise<void> | null = null;
+
+  /**
+   * Creates a PostgreSQL capability bundle store.
+   *
+   * @param options - Pool ownership, connection, and migration options.
+   */
+  constructor(options: PostgreSQLCapabilityBundleStoreOptions) {
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+
+    this.pool = typeof options.poolConfig === "string"
+      ? new Pool({ connectionString: options.poolConfig })
+      : options.poolConfig
+        ? new Pool(options.poolConfig)
+        : (() => { throw new Error("Either pool or poolConfig must be provided"); })();
+
+    this.migrationManager = new MigrationManager(this.pool);
+    this.migrationManager.register(createCapabilityBundlesTable);
+
+    if (options.autoMigrate !== false) {
+      this.startInitialization();
+    }
+  }
+
+  /**
+   * Applies pending migrations for an internally managed pool.
+   *
+   * @returns A shared promise that resolves when initialization completes.
+   */
+  async initialize(): Promise<void> {
+    if (this.initialized) return;
+    if (this.initPromise) return this.initPromise;
+    return this.startInitialization();
+  }
+
+  private startInitialization(): Promise<void> {
+    this.initPromise = this.migrationManager.migrate().then(() => {
+      this.initialized = true;
+    });
+    void this.initPromise.catch(() => undefined);
+    return this.initPromise;
+  }
+
+  /**
+   * Closes the pool when it was created by this store.
+   *
+   * @returns A promise that resolves after owned resources are released.
+   */
+  async dispose(): Promise<void> {
+    if (this.ownsPool) await this.pool.end();
+  }
+
+  private async ready(): Promise<void> {
+    if (!this.initialized) await this.initialize();
+  }
+
+  async listByTenant(tenantId: string): Promise<CapabilityBundle[]> {
+    await this.ready();
+    const result = await this.pool.query<Row>(
+      "SELECT id, tenant_id, bundle_key, name, description, capabilities, created_at, updated_at::text AS updated_at FROM lattice_capability_bundles WHERE tenant_id = $1 ORDER BY created_at",
+      [tenantId],
+    );
+    return result.rows.map(map);
+  }
+
+  async getById(tenantId: string, id: string): Promise<CapabilityBundle | null> {
+    await this.ready();
+    const result = await this.pool.query<Row>(
+      "SELECT id, tenant_id, bundle_key, name, description, capabilities, created_at, updated_at::text AS updated_at FROM lattice_capability_bundles WHERE tenant_id = $1 AND id = $2",
+      [tenantId, id],
+    );
+    return result.rows[0] ? map(result.rows[0]) : null;
+  }
+
+  async getManyByIds(tenantId: string, ids: string[]): Promise<CapabilityBundle[]> {
+    await this.ready();
+    if (ids.length === 0) return [];
+    const result = await this.pool.query<Row>(
+      "SELECT id, tenant_id, bundle_key, name, description, capabilities, created_at, updated_at::text AS updated_at FROM lattice_capability_bundles WHERE tenant_id = $1 AND id = ANY($2::uuid[])",
+      [tenantId, ids],
+    );
+    return result.rows.map(map);
+  }
+
+  async create(
+    tenantId: string,
+    input: InternalCreateCapabilityBundleInput,
+  ): Promise<CapabilityBundle> {
+    await this.ready();
+    try {
+      const result = await this.pool.query<Row>(
+        `INSERT INTO lattice_capability_bundles
+         (id, tenant_id, bundle_key, name, description, capabilities)
+         VALUES ($1, $2, $3, $4, $5, $6::jsonb)
+          RETURNING id, tenant_id, bundle_key, name, description, capabilities, created_at, updated_at::text AS updated_at`,
+        [
+          randomUUID(),
+          tenantId,
+          input.key,
+          input.name,
+          input.description ?? null,
+          JSON.stringify(input.capabilities),
+        ],
+      );
+      return map(result.rows[0]);
+    } catch (error) {
+      if (isDuplicate(error)) throw new Error(duplicateMessage);
+      throw error;
+    }
+  }
+
+  async update(
+    tenantId: string,
+    id: string,
+    input: InternalUpdateCapabilityBundleInput,
+  ): Promise<CapabilityBundle | { status: "conflict" } | null> {
+    await this.ready();
+    const fields: string[] = [];
+    const values: unknown[] = [];
+    const add = (field: string, value: unknown): void => {
+      values.push(value);
+      fields.push(`${field} = $${values.length}`);
+    };
+    if (input.name !== undefined) add("name", input.name);
+    if (Object.prototype.hasOwnProperty.call(input, "description")) {
+      add("description", input.description ?? null);
+    }
+    if (input.capabilities !== undefined) {
+      add("capabilities", JSON.stringify(input.capabilities));
+    }
+    if (fields.length === 0) {
+      const expectedUpdatedAt = input.expectedUpdatedAt;
+      const result = await this.pool.query<Row>(
+        `SELECT id, tenant_id, bundle_key, name, description, capabilities, created_at, updated_at::text AS updated_at
+         FROM lattice_capability_bundles
+         WHERE tenant_id = $1 AND id = $2${expectedUpdatedAt === undefined ? "" : " AND updated_at::text = $3"}`,
+        expectedUpdatedAt === undefined ? [tenantId, id] : [tenantId, id, expectedUpdatedAt],
+      );
+      if (result.rows[0]) return map(result.rows[0]);
+      return expectedUpdatedAt === undefined ? null : { status: "conflict" };
+    }
+
+     const expectedUpdatedAt = input.expectedUpdatedAt;
+     values.push(tenantId, id);
+     if (expectedUpdatedAt !== undefined) values.push(expectedUpdatedAt);
+     const tenantParam = values.length - (expectedUpdatedAt === undefined ? 1 : 2);
+     const idParam = tenantParam + 1;
+     const revisionPredicate = expectedUpdatedAt === undefined ? "" : ` AND updated_at::text = $${values.length}`;
+    try {
+      const result = await this.pool.query<Row>(
+         `UPDATE lattice_capability_bundles
+          SET ${fields.join(", ")}, updated_at = GREATEST(updated_at + interval '1 microsecond', clock_timestamp())
+          WHERE tenant_id = $${tenantParam} AND id = $${idParam}${revisionPredicate}
+         RETURNING id, tenant_id, bundle_key, name, description, capabilities, created_at, updated_at::text AS updated_at`,
+        values,
+      );
+       if (result.rows[0]) return map(result.rows[0]);
+       return expectedUpdatedAt === undefined ? null : { status: "conflict" };
+    } catch (error) {
+      if (isDuplicate(error)) throw new Error(duplicateMessage);
+      throw error;
+    }
+  }
+
+  async deleteIfUnreferenced(tenantId: string, id: string): Promise<CapabilityBundleDeleteResult> {
+    await this.ready();
+    const client = await this.pool.connect();
+    try {
+      await client.query("BEGIN");
+      await client.query(
+        "SELECT pg_advisory_xact_lock(hashtextextended($1 || ':project-mutations', 0))",
+        [tenantId],
+      );
+      const projectIds = await client.query<{ id: string }>(
+        `SELECT id FROM lattice_projects
+         WHERE tenant_id = $1
+         ORDER BY id`,
+        [tenantId],
+      );
+      if (projectIds.rows.length > 0) {
+        await client.query(
+          `SELECT pg_advisory_xact_lock(hashtextextended($1 || ':project:' || project_id, 0))
+           FROM unnest($2::text[]) AS project_id
+           ORDER BY project_id`,
+          [tenantId, projectIds.rows.map((project) => project.id)],
+        );
+      }
+      await client.query(
+        `SELECT id FROM lattice_projects
+         WHERE tenant_id = $1
+         FOR UPDATE`,
+        [tenantId],
+      );
+      await client.query("SELECT pg_advisory_xact_lock(hashtextextended($1 || ':' || $2, 0))", [tenantId, id]);
+      const referenced = await client.query(
+        `SELECT 1 FROM lattice_projects
+         WHERE tenant_id = $1
+         AND jsonb_typeof(config->'capabilityBundleIds') = 'array'
+         AND COALESCE(config->'capabilityBundleIds', '[]'::jsonb) ? $2
+         LIMIT 1`,
+        [tenantId, id],
+      );
+      if (referenced.rows.length > 0) {
+        await client.query("COMMIT");
+        return "in_use";
+      }
+      const deleted = await client.query(
+        "DELETE FROM lattice_capability_bundles WHERE tenant_id = $1 AND id = $2 RETURNING id",
+        [tenantId, id],
+      );
+      await client.query("COMMIT");
+      return deleted.rows.length > 0 ? "deleted" : "not_found";
+    } catch (error) {
+      await client.query("ROLLBACK");
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+}