@axiom-lattice/pg-stores 3.1.1 → 3.1.2

package/src/stores/ThreadMessageQueueStore.ts

@@ -13,12 +13,19 @@ import type {
   ThreadInfo,
   QueueScope,
 } from "@axiom-lattice/core";
+import {
+  parseQueuedExecutionMode,
+  parseTrustedRunContext,
+  type QueuedExecutionMode,
+  type TrustedRunContext,
+} from "@axiom-lattice/protocols";
 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";
 import { addWorkspaceProjectToQueue } from "../migrations/add_workspace_project_to_queue";
+import { addTrustedRunContextColumn } from "../migrations/add_trusted_run_context_column";
 
 export type { PendingMessage, AddMessageParams, ThreadInfo };
 
@@ -59,6 +66,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
     this.migrationManager.register(addCustomRunConfigColumn);
     this.migrationManager.register(alterMessageQueueIdColumn);
     this.migrationManager.register(addWorkspaceProjectToQueue);
+    this.migrationManager.register(addTrustedRunContextColumn);
 
     if (options.autoMigrate !== false) {
       this.initialize().catch((error) => {
@@ -93,6 +101,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
    * Add message to queue
    */
   async addMessage(params: AddMessageParams): Promise<PendingMessage> {
+    const trusted = validateQueueTrust(params);
     const { threadId, tenantId, assistantId, workspaceId, projectId, content, type = "human", priority = 0, command, custom_run_config, id } = params;
 
     // Get current max sequence (no lock needed without unique constraint)
@@ -106,16 +115,17 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
 
     const result = await this.pool.query(
       `INSERT INTO lattice_thread_message_queue 
-       (id, thread_id, tenant_id, assistant_id, workspace_id, project_id, message_content, message_type, sequence_order, priority, command, custom_run_config)
-       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
+       (id, thread_id, tenant_id, assistant_id, workspace_id, project_id, message_content, message_type, sequence_order, priority, command, custom_run_config, trusted_run_context, execution_mode)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
        RETURNING *`,
-      [id || crypto.randomUUID(), threadId, tenantId, assistantId, workspaceId ?? null, projectId ?? null, JSON.stringify(content), type, nextSeq, priority, command ? JSON.stringify(command) : null, custom_run_config ? JSON.stringify(custom_run_config) : null]
+      [id || crypto.randomUUID(), threadId, tenantId, assistantId, workspaceId ?? null, projectId ?? null, JSON.stringify(content), type, nextSeq, priority, command ? JSON.stringify(command) : null, custom_run_config ? JSON.stringify(custom_run_config) : null, trusted.trustedRunContext ? JSON.stringify(trusted.trustedRunContext) : null, trusted.executionMode ?? null]
     );
 
     return this.rowToMessage(result.rows[0]);
   }
 
   async addMessageIfCapacity(params: AddMessageParams, maxSize: number): Promise<boolean> {
+    const trusted = validateQueueTrust(params);
     if (maxSize === Infinity) {
       await this.addMessage(params);
       return true;
@@ -140,9 +150,9 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
         [params.threadId],
       );
       const result = await client.query(
-        `INSERT INTO lattice_thread_message_queue (id, thread_id, tenant_id, assistant_id, workspace_id, project_id, message_content, message_type, sequence_order, priority, command, custom_run_config)
-         VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) RETURNING *`,
-        [params.id || crypto.randomUUID(), params.threadId, params.tenantId, params.assistantId, params.workspaceId ?? null, params.projectId ?? null, JSON.stringify(params.content), params.type || "human", seq.rows[0].next_seq, params.priority ?? 0, params.command ? JSON.stringify(params.command) : null, params.custom_run_config ? JSON.stringify(params.custom_run_config) : null],
+        `INSERT INTO lattice_thread_message_queue (id, thread_id, tenant_id, assistant_id, workspace_id, project_id, message_content, message_type, sequence_order, priority, command, custom_run_config, trusted_run_context, execution_mode)
+         VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) RETURNING *`,
+        [params.id || crypto.randomUUID(), params.threadId, params.tenantId, params.assistantId, params.workspaceId ?? null, params.projectId ?? null, JSON.stringify(params.content), params.type || "human", seq.rows[0].next_seq, params.priority ?? 0, params.command ? JSON.stringify(params.command) : null, params.custom_run_config ? JSON.stringify(params.custom_run_config) : null, trusted.trustedRunContext ? JSON.stringify(trusted.trustedRunContext) : null, trusted.executionMode ?? null],
       );
       await client.query("COMMIT");
       return Boolean(result.rows[0]);
@@ -159,6 +169,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
    * Uses priority=100 to ensure message is processed first
    */
   async addMessageAtHead(params: AddMessageParams): Promise<PendingMessage> {
+    const trusted = validateQueueTrust(params);
     const { threadId, tenantId, assistantId, workspaceId, projectId, content, type = "human", command, custom_run_config, id } = params;
 
     const resolvedTenantId = tenantId!;
@@ -176,10 +187,10 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
     // Insert with high priority (100)
     const result = await this.pool.query(
       `INSERT INTO lattice_thread_message_queue 
-       (id, thread_id, tenant_id, assistant_id, workspace_id, project_id, message_content, message_type, sequence_order, priority, command, custom_run_config)
-       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, 100, $10, $11)
+       (id, thread_id, tenant_id, assistant_id, workspace_id, project_id, message_content, message_type, sequence_order, priority, command, custom_run_config, trusted_run_context, execution_mode)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, 100, $10, $11, $12, $13)
        RETURNING *`,
-      [id || crypto.randomUUID(), threadId, resolvedTenantId, resolvedAssistantId, workspaceId ?? null, projectId ?? null, JSON.stringify(content), type, nextSeq, command ? JSON.stringify(command) : null, custom_run_config ? JSON.stringify(custom_run_config) : null]
+      [id || crypto.randomUUID(), threadId, resolvedTenantId, resolvedAssistantId, workspaceId ?? null, projectId ?? null, JSON.stringify(content), type, nextSeq, command ? JSON.stringify(command) : null, custom_run_config ? JSON.stringify(custom_run_config) : null, trusted.trustedRunContext ? JSON.stringify(trusted.trustedRunContext) : null, trusted.executionMode ?? null]
     );
 
     return this.rowToMessage(result.rows[0]);
@@ -322,10 +333,28 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
       custom_run_config: row.custom_run_config
         ? (typeof row.custom_run_config === "string" ? JSON.parse(row.custom_run_config) : row.custom_run_config)
         : undefined,
+      trusted_run_context: row.trusted_run_context
+        ? parseTrustedRunContext(typeof row.trusted_run_context === "string" ? JSON.parse(row.trusted_run_context) : row.trusted_run_context)
+        : undefined,
+      execution_mode: row.execution_mode == null ? undefined : parseQueuedExecutionMode(row.execution_mode),
     };
   }
 }
 
+function validateQueueTrust(params: AddMessageParams): {
+  trustedRunContext: TrustedRunContext | undefined;
+  executionMode: QueuedExecutionMode | undefined;
+} {
+  return {
+    trustedRunContext: params.trusted_run_context === undefined
+      ? undefined
+      : parseTrustedRunContext(params.trusted_run_context),
+    executionMode: params.execution_mode === undefined
+      ? undefined
+      : parseQueuedExecutionMode(params.execution_mode),
+  };
+}
+
 function scopeClause(scope: QueueScope | undefined, start: number): { sql: string; params: unknown[] } {
   if (!scope) return { sql: "", params: [] };
   const params: unknown[] = [];