@axiom-lattice/pg-stores 3.2.1 → 3.3.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "3.2.1",
+  "version": "3.3.0",
   "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/protocols": "4.6.0",
-    "@axiom-lattice/core": "6.1.1"
+    "@axiom-lattice/core": "6.3.0",
+    "@axiom-lattice/protocols": "4.7.0"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",

package/src/__tests__/PostgreSQLLeaseStore.test.ts

@@ -0,0 +1,43 @@
+import { randomUUID } from "node:crypto";
+import { PostgreSQLLeaseStore } from "../stores/PostgreSQLLeaseStore";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const maybeDescribe = TEST_DATABASE_URL ? describe : describe.skip;
+
+maybeDescribe("PostgreSQLLeaseStore", () => {
+  it("grants exactly one of two concurrent acquires", async () => {
+    const store = new PostgreSQLLeaseStore({ poolConfig: TEST_DATABASE_URL! });
+    await store.initialize();
+    const key = `test-${randomUUID()}`;
+    const [a, b] = await Promise.all([
+      store.tryAcquire(key, "a", 5000),
+      store.tryAcquire(key, "b", 5000),
+    ]);
+    expect([a, b].filter(Boolean)).toHaveLength(1);
+    expect((await store.get(key))?.holder).toBe(a ? "a" : "b");
+    await store.release(key, a ? "a" : "b");
+    await store.dispose();
+  });
+
+  it("lets another holder take over an expired lease", async () => {
+    const store = new PostgreSQLLeaseStore({ poolConfig: TEST_DATABASE_URL! });
+    await store.initialize();
+    const key = `test-${randomUUID()}`;
+    expect(await store.tryAcquire(key, "a", 1)).toBe(true);
+    await new Promise((r) => setTimeout(r, 50));
+    expect(await store.tryAcquire(key, "b", 5000)).toBe(true);
+    expect((await store.get(key))?.holder).toBe("b");
+    await store.dispose();
+  });
+
+  it("is idempotent for the same holder re-acquiring its own live lease", async () => {
+    const store = new PostgreSQLLeaseStore({ poolConfig: TEST_DATABASE_URL! });
+    await store.initialize();
+    const key = `test-${randomUUID()}`;
+    expect(await store.tryAcquire(key, "a", 5000)).toBe(true);
+    const first = await store.get(key);
+    expect(await store.tryAcquire(key, "a", 5000)).toBe(true);
+    expect((await store.get(key))?.acquiredAt).toBe(first?.acquiredAt);
+    await store.dispose();
+  });
+});

package/src/__tests__/PostgreSQLOutboxStore.test.ts

@@ -0,0 +1,77 @@
+import { Pool } from "pg";
+import { PostgreSQLOutboxStore } from "../stores/PostgreSQLOutboxStore";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const maybeDescribe = TEST_DATABASE_URL ? describe : describe.skip;
+
+maybeDescribe("PostgreSQLOutboxStore", () => {
+  it("appends events and pages them with a cursor", async () => {
+    const store = new PostgreSQLOutboxStore({ poolConfig: TEST_DATABASE_URL! });
+    await store.initialize();
+
+    const before = await store.latestId();
+    const first = await store.append({ topic: "t", payload: { n: 1 }, originReplicaId: "a" });
+    const second = await store.append({ topic: "t", payload: { n: 2 }, originReplicaId: "b" });
+
+    const appended = await store.listAfter(before, 10);
+    expect(appended.map((e) => e.id)).toEqual([first.id, second.id]);
+    expect(Number(first.id)).toBeLessThan(Number(second.id));
+
+    const all = await store.listAfter(null, 1000);
+    const allIds = all.map((e) => e.id);
+    const numeric = allIds.map(Number);
+    expect(numeric).toEqual([...numeric].sort((a, b) => a - b));
+    expect(allIds).toEqual(expect.arrayContaining([first.id, second.id]));
+    expect(allIds.indexOf(first.id)).toBeLessThan(allIds.indexOf(second.id));
+
+    const afterFirst = await store.listAfter(first.id, 1000);
+    expect(afterFirst.every((e) => Number(e.id) > Number(first.id))).toBe(true);
+    expect(afterFirst.map((e) => e.id)).toContain(second.id);
+
+    expect(await store.latestId()).toBe(second.id);
+
+    await store.dispose();
+  });
+
+  it("assigns distinct, increasing ids to concurrent appends", async () => {
+    const store = new PostgreSQLOutboxStore({ poolConfig: TEST_DATABASE_URL! });
+    await store.initialize();
+
+    const events = await Promise.all(
+      Array.from({ length: 5 }, (_, i) =>
+        store.append({ topic: "t", payload: { i }, originReplicaId: "c" }),
+      ),
+    );
+    const ids = events.map((e) => Number(e.id)).sort((a, b) => a - b);
+    expect(new Set(ids).size).toBe(5);
+    for (let i = 1; i < ids.length; i++) {
+      expect(ids[i]).toBeGreaterThan(ids[i - 1]);
+    }
+
+    await store.dispose();
+  });
+
+  it("prunes rows older than the threshold and keeps recent ones", async () => {
+    const store = new PostgreSQLOutboxStore({ poolConfig: TEST_DATABASE_URL! });
+    await store.initialize();
+    const pool = new Pool({ connectionString: TEST_DATABASE_URL! });
+
+    const recent = await store.append({ topic: "t", payload: { n: 1 }, originReplicaId: "a" });
+    const old = await pool.query(
+      `INSERT INTO lattice_outbox (topic, payload, origin_replica_id, created_at)
+       VALUES ('t', '{}'::jsonb, 'a', now() - interval '2 days')
+       RETURNING id`,
+    );
+    const oldId = String(old.rows[0].id);
+
+    const removed = await store.pruneOlderThan(24 * 60 * 60 * 1000);
+
+    expect(removed).toBeGreaterThanOrEqual(1);
+    const remaining = (await store.listAfter(null, 1000)).map((event) => event.id);
+    expect(remaining).toContain(recent.id);
+    expect(remaining).not.toContain(oldId);
+
+    await pool.end();
+    await store.dispose();
+  });
+});

package/src/__tests__/PostgreSQLScheduleStorage.test.ts

@@ -0,0 +1,162 @@
+import { randomUUID } from "node:crypto";
+
+import {
+  ScheduleExecutionType,
+  ScheduledTaskDefinition,
+  ScheduledTaskStatus,
+} from "@axiom-lattice/protocols";
+
+import { PostgreSQLScheduleStorage } from "../stores/PostgreSQLScheduleStorage";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const maybeDescribe = TEST_DATABASE_URL ? describe : describe.skip;
+
+const buildTask = (
+  taskId: string,
+  overrides: Partial<ScheduledTaskDefinition> = {}
+): ScheduledTaskDefinition => {
+  const now = Date.now();
+  return {
+    taskId,
+    taskType: "test",
+    payload: {},
+    tenantId: "tenant-1",
+    executionType: ScheduleExecutionType.ONCE,
+    executeAt: now - 1000,
+    status: ScheduledTaskStatus.PENDING,
+    runCount: 0,
+    retryCount: 0,
+    maxRetries: 0,
+    createdAt: now,
+    updatedAt: now,
+    ...overrides,
+  };
+};
+
+maybeDescribe("PostgreSQLScheduleStorage.claimScheduledRun", () => {
+  it("claims a pending task exactly once", async () => {
+    const store = new PostgreSQLScheduleStorage({
+      poolConfig: TEST_DATABASE_URL!,
+    });
+    await store.initialize();
+    const taskId = `test-${randomUUID()}`;
+    await store.save(buildTask(taskId));
+
+    expect(await store.claimScheduledRun(taskId)).toBe(true);
+    expect(await store.claimScheduledRun(taskId)).toBe(false);
+
+    const stored = await store.get(taskId);
+    expect(stored?.status).toBe(ScheduledTaskStatus.RUNNING);
+
+    await store.delete(taskId);
+    await store.dispose();
+  });
+
+  it("grants exactly one of two concurrent claims", async () => {
+    const store = new PostgreSQLScheduleStorage({
+      poolConfig: TEST_DATABASE_URL!,
+    });
+    await store.initialize();
+    const taskId = `test-${randomUUID()}`;
+    await store.save(buildTask(taskId));
+
+    const [a, b] = await Promise.all([
+      store.claimScheduledRun(taskId),
+      store.claimScheduledRun(taskId),
+    ]);
+    expect([a, b].filter(Boolean)).toHaveLength(1);
+
+    await store.delete(taskId);
+    await store.dispose();
+  });
+
+  it("returns false for an unknown task", async () => {
+    const store = new PostgreSQLScheduleStorage({
+      poolConfig: TEST_DATABASE_URL!,
+    });
+    await store.initialize();
+
+    expect(await store.claimScheduledRun(`missing-${randomUUID()}`)).toBe(false);
+
+    await store.dispose();
+  });
+});
+
+maybeDescribe("PostgreSQLScheduleStorage.requeueStaleRunning", () => {
+  it("requeues a running task older than the threshold", async () => {
+    const store = new PostgreSQLScheduleStorage({
+      poolConfig: TEST_DATABASE_URL!,
+    });
+    await store.initialize();
+    const taskId = `test-${randomUUID()}`;
+    await store.save(
+      buildTask(taskId, {
+        status: ScheduledTaskStatus.RUNNING,
+        updatedAt: Date.now() - 60_000,
+      })
+    );
+
+    expect(await store.requeueStaleRunning(30_000)).toBe(1);
+
+    const stored = await store.get(taskId);
+    expect(stored?.status).toBe(ScheduledTaskStatus.PENDING);
+
+    await store.delete(taskId);
+    await store.dispose();
+  });
+
+  it("leaves a recently updated running task alone", async () => {
+    const store = new PostgreSQLScheduleStorage({
+      poolConfig: TEST_DATABASE_URL!,
+    });
+    await store.initialize();
+    const taskId = `test-${randomUUID()}`;
+    await store.save(
+      buildTask(taskId, {
+        status: ScheduledTaskStatus.RUNNING,
+        updatedAt: Date.now() - 1_000,
+      })
+    );
+
+    expect(await store.requeueStaleRunning(30_000)).toBe(0);
+
+    const stored = await store.get(taskId);
+    expect(stored?.status).toBe(ScheduledTaskStatus.RUNNING);
+
+    await store.delete(taskId);
+    await store.dispose();
+  });
+
+  it("never touches pending or terminal tasks", async () => {
+    const store = new PostgreSQLScheduleStorage({
+      poolConfig: TEST_DATABASE_URL!,
+    });
+    await store.initialize();
+    const pendingId = `test-${randomUUID()}`;
+    const completedId = `test-${randomUUID()}`;
+    await store.save(
+      buildTask(pendingId, {
+        status: ScheduledTaskStatus.PENDING,
+        updatedAt: Date.now() - 60_000,
+      })
+    );
+    await store.save(
+      buildTask(completedId, {
+        status: ScheduledTaskStatus.COMPLETED,
+        updatedAt: Date.now() - 60_000,
+      })
+    );
+
+    expect(await store.requeueStaleRunning(30_000)).toBe(0);
+    expect((await store.get(pendingId))?.status).toBe(
+      ScheduledTaskStatus.PENDING
+    );
+    expect((await store.get(completedId))?.status).toBe(
+      ScheduledTaskStatus.COMPLETED
+    );
+
+    await store.delete(pendingId);
+    await store.delete(completedId);
+    await store.dispose();
+  });
+});

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

@@ -0,0 +1,82 @@
+import { randomUUID } from "node:crypto";
+import { ThreadMessageQueueStore } from "../stores/ThreadMessageQueueStore";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const maybeDescribe = TEST_DATABASE_URL ? describe : describe.skip;
+
+maybeDescribe("ThreadMessageQueueStore claim/lease (PostgreSQL)", () => {
+  let store: ThreadMessageQueueStore;
+  const threadIds: string[] = [];
+
+  function newThreadId(): string {
+    const threadId = `claim-${randomUUID()}`;
+    threadIds.push(threadId);
+    return threadId;
+  }
+
+  beforeAll(async () => {
+    store = new ThreadMessageQueueStore({ poolConfig: TEST_DATABASE_URL!, autoMigrate: true });
+    await store.initialize();
+  });
+
+  afterAll(async () => {
+    for (const threadId of threadIds) {
+      await store.clearMessages(threadId);
+    }
+    await store.dispose();
+  });
+
+  it("grants disjoint sets to concurrent claimants", async () => {
+    const threadId = newThreadId();
+    for (let i = 0; i < 4; i += 1) {
+      await store.addMessage({
+        threadId,
+        tenantId: "tenant",
+        assistantId: "assistant",
+        content: { id: randomUUID(), message: `message-${i}`, queueMessage: { input: {} } },
+        id: randomUUID(),
+      });
+    }
+
+    const [a, b] = await Promise.all([
+      store.claimPendingMessages(threadId, undefined, { owner: "r1", ttlMs: 5000, limit: 2 }),
+      store.claimPendingMessages(threadId, undefined, { owner: "r2", ttlMs: 5000, limit: 2 }),
+    ]);
+
+    const ids = [...a, ...b].map((m) => m.id);
+    expect(ids).toHaveLength(4);
+    expect(new Set(ids).size).toBe(4);
+    expect(await store.getPendingMessages(threadId)).toHaveLength(0);
+    expect(await store.getProcessingMessages(threadId)).toHaveLength(4);
+  });
+
+  it("reclaims only expired processing rows", async () => {
+    const liveThread = newThreadId();
+    const expiredThread = newThreadId();
+    await store.addMessage({ threadId: liveThread, tenantId: "tenant", assistantId: "assistant", content: { id: randomUUID(), message: "live", queueMessage: { input: {} } }, id: randomUUID() });
+    await store.addMessage({ threadId: expiredThread, tenantId: "tenant", assistantId: "assistant", content: { id: randomUUID(), message: "expired", queueMessage: { input: {} } }, id: randomUUID() });
+
+    await store.claimPendingMessages(liveThread, undefined, { owner: "r1", ttlMs: 60000 });
+    await store.claimPendingMessages(expiredThread, undefined, { owner: "r1", ttlMs: 1 });
+    await new Promise((resolve) => setTimeout(resolve, 50));
+
+    expect(await store.reclaimExpiredProcessing(liveThread)).toBe(0);
+    expect(await store.reclaimExpiredProcessing(expiredThread)).toBe(1);
+    expect(await store.getPendingMessages(expiredThread)).toHaveLength(1);
+    expect(await store.getPendingMessages(liveThread)).toHaveLength(0);
+  });
+
+  it("round-trips a bare-string message content", async () => {
+    const threadId = newThreadId();
+    await store.addMessage({
+      threadId,
+      tenantId: "ten",
+      assistantId: "a",
+      content: "hello",
+      id: randomUUID(),
+    });
+
+    const [pending] = await store.getPendingMessages(threadId);
+    expect(pending.content).toBe("hello");
+  });
+});

package/src/__tests__/ThreadMessageQueueStore.test.ts

@@ -46,7 +46,7 @@ describe('ThreadMessageQueueStore', () => {
       'persists trusted metadata through %s',
       async (method) => {
         const row = {
-          id: 'input-1', message_content: JSON.stringify('message'), message_type: 'human',
+          id: 'input-1', message_content: 'message', message_type: 'human',
           sequence_order: 1, created_at: new Date(), trusted_run_context: trusted,
           execution_mode: 'followup',
         };
@@ -144,7 +144,7 @@ describe('ThreadMessageQueueStore', () => {
         thread_id: mockThreadId,
         tenant_id: mockTenantId,
         assistant_id: mockAssistantId,
-        message_content: JSON.stringify('test content'),
+        message_content: 'test content',
         message_type: 'human',
         sequence_order: 1,
         status: 'pending',
@@ -176,7 +176,7 @@ describe('ThreadMessageQueueStore', () => {
           rows: [{ 
             id: 'msg-5',
             sequence_order: 5,
-            message_content: JSON.stringify('content'),
+            message_content: 'content',
           }] 
         });
 
@@ -205,7 +205,7 @@ describe('ThreadMessageQueueStore', () => {
             thread_id: mockThreadId,
             tenant_id: mockTenantId,
             assistant_id: mockAssistantId,
-            message_content: JSON.stringify('test content'),
+            message_content: 'test content',
             message_type: 'human',
             sequence_order: 1,
             status: 'pending',
@@ -241,7 +241,7 @@ describe('ThreadMessageQueueStore', () => {
             id: 'msg-head',
             sequence_order: 1,
             priority: 100,
-            message_content: JSON.stringify('head content'),
+            message_content: 'head content',
           }] 
         }); // Insert
 
@@ -267,7 +267,7 @@ describe('ThreadMessageQueueStore', () => {
             id: 'msg-head',
             sequence_order: 0,
             priority: 100,
-            message_content: JSON.stringify('head content'),
+            message_content: 'head content',
           }] 
         });
 
@@ -288,14 +288,14 @@ describe('ThreadMessageQueueStore', () => {
       const mockMessages = [
         {
           id: 'msg-1',
-          message_content: JSON.stringify('first'),
+          message_content: 'first',
           message_type: 'human',
           sequence_order: 1,
           created_at: new Date('2024-01-01'),
         },
         {
           id: 'msg-2',
-          message_content: JSON.stringify('second'),
+          message_content: 'second',
           message_type: 'human',
           sequence_order: 2,
           created_at: new Date('2024-01-02'),
@@ -335,7 +335,7 @@ describe('ThreadMessageQueueStore', () => {
         rows: [
           {
             id: 'msg-1',
-            message_content: JSON.stringify({
+            message_content: {
               id: 'msg-1',
               message: 'Clarification answers submitted',
               queueMessage: {
@@ -344,7 +344,7 @@ describe('ThreadMessageQueueStore', () => {
                   message: 'Clarification answers submitted',
                 },
               },
-            }),
+            },
             command: JSON.stringify(command),
             custom_run_config: JSON.stringify(customRunConfig),
             message_type: 'human',
@@ -364,7 +364,7 @@ describe('ThreadMessageQueueStore', () => {
 
     it('hydrates trusted metadata for pending and processing rows', async () => {
       const row = {
-        id: 'input-1', message_content: JSON.stringify('message'), message_type: 'human',
+        id: 'input-1', message_content: 'message', message_type: 'human',
         sequence_order: 1, created_at: new Date(), trusted_run_context: JSON.stringify(trusted),
         execution_mode: 'followup',
       };
@@ -380,7 +380,7 @@ describe('ThreadMessageQueueStore', () => {
 
     it('rejects hostile trusted JSON read from PostgreSQL', async () => {
       mockQuery.mockResolvedValueOnce({ rows: [{
-        id: 'hostile', message_content: JSON.stringify('message'), message_type: 'human',
+        id: 'hostile', message_content: 'message', message_type: 'human',
         sequence_order: 1, created_at: new Date(),
         trusted_run_context: { projectRoom: { tenantId: mockTenantId } }, execution_mode: 'followup',
       }] });
@@ -547,4 +547,63 @@ describe('ThreadMessageQueueStore', () => {
     });
   });
 
+  describe('claimPendingMessages', () => {
+    it('claims with lease columns bounded by an ordered SKIP LOCKED subquery', async () => {
+      const row = {
+        id: 'msg-1', message_content: 'content', message_type: 'human',
+        sequence_order: 1, created_at: new Date(),
+      };
+      mockQuery.mockResolvedValueOnce({ rows: [row] });
+
+      const claimed = await store.claimPendingMessages(mockThreadId, undefined, { owner: 'r1', ttlMs: 5000, limit: 2 });
+
+      expect(claimed.map((message) => message.id)).toEqual(['msg-1']);
+      const sql = mockQuery.mock.calls[0][0] as string;
+      expect(sql).toContain("status = 'processing'");
+      expect(sql).toContain('locked_by = $2');
+      expect(sql).toContain("lease_expires_at = now() + ($3::bigint || ' milliseconds')::interval");
+      expect(sql).toContain('ORDER BY priority DESC, sequence_order ASC LIMIT 2');
+      expect(sql).toContain('FOR UPDATE SKIP LOCKED');
+      expect(mockQuery.mock.calls[0][1]).toEqual([mockThreadId, 'r1', 5000]);
+    });
+
+    it('applies scope predicates without colliding with the owner parameter', async () => {
+      mockQuery.mockResolvedValueOnce({ rows: [] });
+
+      await store.claimPendingMessages(mockThreadId, {
+        tenantId: mockTenantId, assistantId: mockAssistantId, workspaceId: undefined, projectId: undefined,
+      }, { owner: 'r1', ttlMs: 5000 });
+
+      const [sql, params] = mockQuery.mock.calls[0] as [string, unknown[]];
+      expect(sql).toMatch(/tenant_id = \$4.*assistant_id = \$5.*workspace_id IS NULL.*project_id IS NULL/s);
+      expect(params).toEqual([mockThreadId, 'r1', 5000, mockTenantId, mockAssistantId]);
+    });
+  });
+
+  describe('reclaimExpiredProcessing', () => {
+    it('resets only expired leases and returns the row count', async () => {
+      mockQuery.mockResolvedValueOnce({ rowCount: 2 });
+
+      await expect(store.reclaimExpiredProcessing(mockThreadId)).resolves.toBe(2);
+
+      const sql = mockQuery.mock.calls[0][0] as string;
+      expect(sql).toContain("SET status = 'pending', locked_by = NULL, locked_at = NULL, lease_expires_at = NULL");
+      expect(sql).toContain('lease_expires_at <= now()');
+      expect(mockQuery.mock.calls[0][1]).toEqual([mockThreadId]);
+    });
+  });
+
+  describe('resetProcessingToPending', () => {
+    it('clears the lease columns for reset rows', async () => {
+      mockQuery.mockResolvedValueOnce({ rowCount: 1 });
+
+      await store.resetProcessingToPending(mockThreadId);
+
+      const sql = mockQuery.mock.calls[0][0] as string;
+      expect(sql).toContain('locked_by = NULL');
+      expect(sql).toContain('locked_at = NULL');
+      expect(sql).toContain('lease_expires_at = NULL');
+    });
+  });
+
 });

package/src/__tests__/queue_migrations.test.ts

@@ -0,0 +1,18 @@
+import { registerThreadMessageQueueMigrations } from "../migrations/queue_migrations";
+
+describe("registerThreadMessageQueueMigrations", () => {
+  it("registers every thread-message-queue migration, including the claim lease", () => {
+    const names: string[] = [];
+    const fake = {
+      register: (migration: { name: string }): void => {
+        names.push(migration.name);
+      },
+    };
+    registerThreadMessageQueueMigrations(fake as never);
+    // The shared list must stay complete; bump this count when adding a queue
+    // migration (and register it in the shared helper, never inline).
+    expect(names).toHaveLength(7);
+    expect(names).toContain("create_lattice_thread_message_queue_table");
+    expect(names).toContain("add_queue_claim_lease_columns");
+  });
+});

package/src/createPgStoreConfig.ts

@@ -58,12 +58,7 @@ import { createDatabaseConfigsTable } from "./migrations/database_config_migrati
 import { createConnectionConfigsTable } from "./migrations/connection_config_migrations";
 import { createMetricsConfigsTable } from "./migrations/metrics_config_migrations";
 import { createMcpServerConfigsTable } from "./migrations/mcp_server_config_migrations";
-import { createThreadMessageQueueTable } from "./migrations/thread_message_queue_migrations";
-import { addPriorityAndCommandColumns } from "./migrations/add_priority_command_columns";
-import { addCustomRunConfigColumn } from "./migrations/add_custom_run_config_column";
-import { alterMessageQueueIdColumn } from "./migrations/alter_message_queue_id_column";
-import { addWorkspaceProjectToQueue } from "./migrations/add_workspace_project_to_queue";
-import { addTrustedRunContextColumn } from "./migrations/add_trusted_run_context_column";
+import { registerThreadMessageQueueMigrations } from "./migrations/queue_migrations";
 import { createWorkflowTrackingTables, addStepThreadId, addWorkflowRunsTenantStatusUpdatedIndex } from "./migrations/workflow_tracking_migrations";
 import { evalMigrations } from "./migrations/eval_migrations";
 import { createA2AApiKeysTable } from "./migrations/a2a_api_key_migration";
@@ -90,6 +85,8 @@ import { createCapabilityBundlesTable } from "./migrations/capability_bundle_mig
 import { createProjectRoomReadStates, createProjectRoomTables } from "./migrations/project_room_migration";
 import { addProjectRoomMessageRetentionIndex } from "./migrations/project_room_message_retention_migration";
 import { createUserPushSubscriptions } from "./migrations/user_push_subscription_migration";
+import { createDistributedLeases } from "./migrations/distributed_lease_migration";
+import { createOutboxTable } from "./migrations/outbox_migration";
 import {
   createModelProvidersTable,
   createModelProviderModelsTable,
@@ -104,6 +101,8 @@ import { PostgreSQLProjectRoomMessageStore } from "./stores/PostgreSQLProjectRoo
 import { PostgreSQLProjectRoomReadStateStore } from "./stores/PostgreSQLProjectRoomReadStateStore";
 import { PostgreSQLUserPushSubscriptionStore } from "./stores/PostgreSQLUserPushSubscriptionStore";
 import { PostgreSQLModelProviderStore } from "./stores/PostgreSQLModelProviderStore";
+import { PostgreSQLLeaseStore } from "./stores/PostgreSQLLeaseStore";
+import { PostgreSQLOutboxStore } from "./stores/PostgreSQLOutboxStore";
 
 export async function createPgStoreConfig(connectionString: string) {
   const pool = new Pool({ connectionString });
@@ -134,9 +133,7 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(createChannelInstallationsTable);    // v21
   mm.register(addThreadTenantId);           // v23
   mm.register(addProjectConfigColumn);      // v24
-  mm.register(createThreadMessageQueueTable); // v100
-  mm.register(addPriorityAndCommandColumns);  // v101
-  mm.register(alterMessageQueueIdColumn);     // v102
+  registerThreadMessageQueueMigrations(mm);   // v100, v101, v102, v114, v131, v173, v186
   mm.register(createWorkflowTrackingTables);  // v103
   for (const m of evalMigrations) {           // v104-v108
     mm.register(m);
@@ -146,9 +143,7 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(createDatabaseConfigsTable);      // v111
   mm.register(changeAssistantPrimaryKey);       // v112
   mm.register(addStepThreadId);                 // v113
-  mm.register(addCustomRunConfigColumn);        // v114
   mm.register(createA2AApiKeysTable);           // v130
-  mm.register(addWorkspaceProjectToQueue);      // v131
   mm.register(addAssistantOwnerUserId);         // v132
   for (const m of taskMigrations) {             // v133, v139, v140, v141
     mm.register(m);
@@ -172,7 +167,6 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(createCapabilityBundlesTable);   // v170
   mm.register(addTaskWorkItemPendingIndexesMigration); // v171
   mm.register(createProjectRoomTables);        // v172
-  mm.register(addTrustedRunContextColumn);      // v173
   mm.register(addProjectLifecycleEventIndex);   // v175
   mm.register(addAgentWebAppsIdentity);         // v176
   mm.register(createProjectRoomReadStates);     // v177
@@ -182,6 +176,8 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(createModelProviderModelsTable);  // v182
   mm.register(addModelProviderApiStyle);        // v183
   mm.register(addModelProviderModelStale);      // v184
+  mm.register(createDistributedLeases);        // v185
+  mm.register(createOutboxTable);              // v187
 
   await mm.migrate();
 
@@ -228,6 +224,8 @@ export async function createPgStoreConfig(connectionString: string) {
     projectRoomMessage: new PostgreSQLProjectRoomMessageStore(opts),
     projectRoomReadState: new PostgreSQLProjectRoomReadStateStore(opts),
     userPushSubscription: new PostgreSQLUserPushSubscriptionStore(opts),
+    lease: new PostgreSQLLeaseStore(opts),
+    outbox: new PostgreSQLOutboxStore(opts),
     vectorStoreProvider: new PGVectorStoreProvider(pool, connectionString),
     checkpoint,
   };