@axiom-lattice/pg-stores 1.0.34 → 1.0.36

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "1.0.34",
+  "version": "1.0.36",
   "description": "PG stores implementation for Axiom Lattice framework",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -21,8 +21,8 @@
   "license": "MIT",
   "dependencies": {
     "pg": "^8.16.3",
-    "@axiom-lattice/protocols": "2.1.23",
-    "@axiom-lattice/core": "2.1.44"
+    "@axiom-lattice/protocols": "2.1.24",
+    "@axiom-lattice/core": "2.1.46"
   },
   "devDependencies": {
     "@types/node": "^20.11.24",

package/src/__tests__/ThreadMessageQueueStore.test.ts

@@ -89,6 +89,46 @@ describe('ThreadMessageQueueStore', () => {
 
       expect(result.sequence).toBe(5);
     });
+
+    it('should persist custom_run_config in insert query', async () => {
+      const customRunConfig = {
+        modelConfig: {
+          modelKey: 'default',
+        },
+      };
+
+      mockQuery
+        .mockResolvedValueOnce({ rows: [{ next_seq: 1 }] })
+        .mockResolvedValueOnce({
+          rows: [{
+            id: 'msg-1',
+            thread_id: mockThreadId,
+            tenant_id: mockTenantId,
+            assistant_id: mockAssistantId,
+            message_content: JSON.stringify('test content'),
+            message_type: 'human',
+            sequence_order: 1,
+            status: 'pending',
+            created_at: new Date(),
+            custom_run_config: JSON.stringify(customRunConfig),
+          }],
+        });
+
+      await store.addMessage({
+        threadId: mockThreadId,
+        tenantId: mockTenantId,
+        assistantId: mockAssistantId,
+        content: 'test content',
+        type: 'human',
+        custom_run_config: customRunConfig,
+      });
+
+      expect(mockQuery).toHaveBeenNthCalledWith(
+        2,
+        expect.stringContaining('custom_run_config'),
+        expect.arrayContaining([JSON.stringify(customRunConfig)])
+      );
+    });
   });
 
   describe('addMessageAtHead', () => {
@@ -106,7 +146,13 @@ describe('ThreadMessageQueueStore', () => {
           }] 
         }); // Insert
 
-      const result = await store.addMessageAtHead(mockThreadId, 'head content', 'system');
+      const result = await store.addMessageAtHead({
+        threadId: mockThreadId,
+        tenantId: mockTenantId,
+        assistantId: mockAssistantId,
+        content: 'head content',
+        type: 'system',
+      });
 
       expect(result.priority).toBe(100);
       expect(result.sequence).toBe(1);
@@ -127,7 +173,13 @@ describe('ThreadMessageQueueStore', () => {
           }] 
         });
 
-      const result = await store.addMessageAtHead(mockThreadId, 'head content', 'system');
+      const result = await store.addMessageAtHead({
+        threadId: mockThreadId,
+        tenantId: mockTenantId,
+        assistantId: mockAssistantId,
+        content: 'head content',
+        type: 'system',
+      });
 
       expect(result.priority).toBe(100);
     });
@@ -175,6 +227,11 @@ describe('ThreadMessageQueueStore', () => {
           action: 'submit',
         },
       };
+      const customRunConfig = {
+        modelConfig: {
+          modelKey: 'default',
+        },
+      };
 
       mockQuery.mockResolvedValueOnce({
         rows: [
@@ -191,6 +248,7 @@ describe('ThreadMessageQueueStore', () => {
               },
             }),
             command: JSON.stringify(command),
+            custom_run_config: JSON.stringify(customRunConfig),
             message_type: 'human',
             sequence_order: 1,
             created_at: new Date('2024-01-01'),
@@ -202,6 +260,7 @@ describe('ThreadMessageQueueStore', () => {
       const result = await store.getPendingMessages(mockThreadId);
 
       expect(result[0].command).toEqual(command);
+      expect(result[0].custom_run_config).toEqual(customRunConfig);
       expect((result[0].content as { command?: unknown }).command).toBeUndefined();
     });
   });

package/src/migrations/add_custom_run_config_column.ts

@@ -0,0 +1,25 @@
+/**
+ * Add custom_run_config column to thread message queue
+ */
+
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+/**
+ * Migration: Add custom_run_config column
+ */
+export const addCustomRunConfigColumn: Migration = {
+  version: 102,
+  name: "add_custom_run_config_column",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_thread_message_queue
+      ADD COLUMN IF NOT EXISTS custom_run_config JSONB
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query(
+      "ALTER TABLE lattice_thread_message_queue DROP COLUMN IF EXISTS custom_run_config"
+    );
+  },
+};

package/src/stores/PostgreSQLThreadStore.ts

@@ -103,7 +103,11 @@ export class PostgreSQLThreadStore implements ThreadStore {
   /**
    * Get all threads for a specific tenant and assistant
    */
-  async getThreadsByAssistantId(tenantId: string, assistantId: string): Promise<Thread[]> {
+  async getThreadsByAssistantId(
+    tenantId: string,
+    assistantId: string,
+    metadataFilter?: Record<string, string>
+  ): Promise<Thread[]> {
     await this.ensureInitialized();
 
     const result = await this.pool.query<{
@@ -123,7 +127,18 @@ export class PostgreSQLThreadStore implements ThreadStore {
       [tenantId, assistantId]
     );
 
-    return result.rows.map(this.mapRowToThread);
+    let threads = result.rows.map(this.mapRowToThread);
+
+    // Apply metadata filtering if provided
+    if (metadataFilter && Object.keys(metadataFilter).length > 0) {
+      threads = threads.filter((thread) =>
+        Object.entries(metadataFilter).every(
+          ([key, value]) => thread.metadata?.[key] === value
+        )
+      );
+    }
+
+    return threads;
   }
 
   /**

package/src/stores/ThreadMessageQueueStore.ts

@@ -15,6 +15,7 @@ import {
 import { MigrationManager } from "../migrations/migration";
 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";
 
 export { PendingMessage, AddMessageParams, ThreadInfo };
@@ -45,6 +46,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
       const migrationManager = new MigrationManager(pool);
       migrationManager.register(createThreadMessageQueueTable);
       migrationManager.register(addPriorityAndCommandColumns);
+      migrationManager.register(addCustomRunConfigColumn);
       migrationManager.register(alterMessageQueueIdColumn);
       await migrationManager.migrate();
       this.initialized = true;
@@ -63,7 +65,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
    */
   async addMessage(params: AddMessageParams): Promise<PendingMessage> {
     const pool = this.getPool();
-    const { threadId, tenantId, assistantId, content, type = "human", priority = 0, command, id } = params;
+    const { threadId, tenantId, assistantId, content, type = "human", priority = 0, command, custom_run_config, id } = params;
 
     // Get current max sequence (no lock needed without unique constraint)
     const seqResult = await pool.query(
@@ -76,10 +78,10 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
 
     const result = await pool.query(
       `INSERT INTO lattice_thread_message_queue 
-       (id, thread_id, tenant_id, assistant_id, message_content, message_type, sequence_order, priority, command)
-       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
+       (id, thread_id, tenant_id, assistant_id, message_content, message_type, sequence_order, priority, command, custom_run_config)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
        RETURNING *`,
-      [id || crypto.randomUUID(), threadId, tenantId, assistantId, JSON.stringify(content), type, nextSeq, priority, command ? JSON.stringify(command) : null]
+      [id || crypto.randomUUID(), threadId, tenantId, assistantId, JSON.stringify(content), type, nextSeq, priority, command ? JSON.stringify(command) : null, custom_run_config ? JSON.stringify(custom_run_config) : null]
     );
 
     return this.rowToMessage(result.rows[0]);
@@ -89,31 +91,14 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
    * Add message at head of queue (high priority, e.g., STEER/Command messages)
    * Uses priority=100 to ensure message is processed first
    */
-  async addMessageAtHead(
-    threadId: string,
-    content: PendingMessage["content"],
-    type: "human" | "system" = "system",
-    id?: string,
-    command?: PendingMessage["command"]
-  ): Promise<PendingMessage> {
+  async addMessageAtHead(params: AddMessageParams): Promise<PendingMessage> {
     const pool = this.getPool();
+    const { threadId, tenantId, assistantId, content, type = "human", command, custom_run_config, id } = params;
 
-    // Get thread info from existing messages
-    const threadResult = await pool.query(
-      `SELECT tenant_id, assistant_id 
-       FROM lattice_thread_message_queue 
-       WHERE thread_id = $1 
-       LIMIT 1`,
-      [threadId]
-    );
-
-    let tenantId = "default";
-    let assistantId = "";
-
-    if (threadResult.rows.length > 0) {
-      tenantId = threadResult.rows[0].tenant_id;
-      assistantId = threadResult.rows[0].assistant_id;
-    }
+    // Use provided tenantId and assistantId directly
+    // These should always be provided by the caller (Agent)
+    const resolvedTenantId = tenantId!;
+    const resolvedAssistantId = assistantId!;
 
     // Get current max sequence for ordering within same priority
     const seqResult = await pool.query(
@@ -127,10 +112,10 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
     // Insert with high priority (100)
     const result = await pool.query(
       `INSERT INTO lattice_thread_message_queue 
-       (id, thread_id, tenant_id, assistant_id, message_content, message_type, sequence_order, priority, command)
-       VALUES ($1, $2, $3, $4, $5, $6, $7, 100, $8)
+       (id, thread_id, tenant_id, assistant_id, message_content, message_type, sequence_order, priority, command, custom_run_config)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, 100, $8, $9)
        RETURNING *`,
-      [id || crypto.randomUUID(), threadId, tenantId, assistantId, JSON.stringify(content), type, nextSeq, command ? JSON.stringify(command) : null]
+      [id || crypto.randomUUID(), threadId, resolvedTenantId, resolvedAssistantId, JSON.stringify(content), type, nextSeq, command ? JSON.stringify(command) : null, custom_run_config ? JSON.stringify(custom_run_config) : null]
     );
 
     return this.rowToMessage(result.rows[0]);
@@ -284,6 +269,9 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
       createdAt: new Date(row.created_at),
       priority: row.priority || 0,
       command: row.command ? (typeof row.command === "string" ? JSON.parse(row.command) : row.command) : undefined,
+      custom_run_config: row.custom_run_config
+        ? (typeof row.custom_run_config === "string" ? JSON.parse(row.custom_run_config) : row.custom_run_config)
+        : undefined,
     };
   }
 }