@axiom-lattice/pg-stores 3.1.0 → 3.1.2

package/src/stores/PostgreSQLTaskWorkItemStore.ts

@@ -1,16 +1,97 @@
-import type { Pool } from "pg";
-import type { TaskWorkItemStore, TaskWorkItem, CreateWorkItemRequest, CreateWorkItemIfAbsentRequest, TaskWorkItemListFilter } from "@axiom-lattice/protocols";
+import { Pool } from "pg";
+import type { PoolConfig } from "pg";
+import type { TaskWorkItemStore, TaskWorkItem, CreateWorkItemRequest, CreateWorkItemIfAbsentRequest, TaskWorkItemListFilter, ProjectLifecycleEventQuery, TaskMutationSnapshot } from "@axiom-lattice/protocols";
+import { MAX_PENDING_EXECUTION_RESULTS_LIMIT } from "@axiom-lattice/protocols";
 import { v4 } from "uuid";
+import { MigrationManager } from "../migrations/migration";
+import { taskWorkItemMigrations } from "../migrations/task_work_items_migration";
 
+const PROJECT_LIFECYCLE_ACTIONS: ReadonlySet<string> = new Set([
+  "in_progress", "interrupted", "failed", "completed", "cancelled", "reassigned",
+]);
+
+/** PostgreSQL TaskWorkItem store connection and migration options. */
+export interface PostgreSQLTaskWorkItemStoreOptions {
+  /** External caller-owned pool; central migrations must already be applied. */
+  pool?: Pool;
+  /** Connection configuration for a standalone, store-owned pool. */
+  poolConfig?: string | PoolConfig;
+  /** Whether a standalone store eagerly applies migrations. Defaults to true. */
+  autoMigrate?: boolean;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return value !== null && typeof value === "object" && !Array.isArray(value);
+}
+
+function isPool(value: Pool | PostgreSQLTaskWorkItemStoreOptions): value is Pool {
+  return typeof (value as { query?: unknown }).query === "function";
+}
+
+/** PostgreSQL TaskWorkItem persistence with optional standalone migrations. */
 export class PostgreSQLTaskWorkItemStore implements TaskWorkItemStore {
-  constructor(private pool: Pool) {}
+  private pool: Pool;
+  private migrationManager!: MigrationManager;
+  private initialized = false;
+  private ownsPool = false;
+  private initPromise: Promise<void> | null = null;
+
+  constructor(pool: Pool);
+  constructor(options: PostgreSQLTaskWorkItemStoreOptions);
+  constructor(poolOrOptions: Pool | PostgreSQLTaskWorkItemStoreOptions) {
+    if (isPool(poolOrOptions)) {
+      this.pool = poolOrOptions;
+      this.initialized = true;
+      return;
+    }
+    const options = poolOrOptions as PostgreSQLTaskWorkItemStoreOptions;
+    if (options.pool) {
+      this.pool = options.pool;
+      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.ownsPool = true;
+    this.migrationManager = new MigrationManager(this.pool);
+    for (const migration of taskWorkItemMigrations) this.migrationManager.register(migration);
+    if (options.autoMigrate !== false) this.startInitialization();
+  }
+
+  /** Applies the complete standalone TaskWorkItem migration chain once. */
+  async initialize(): Promise<void> {
+    if (this.initialized) return;
+    if (this.initPromise) return this.initPromise;
+    return this.startInitialization();
+  }
+
+  /** Closes the pool only when this store created it from connection configuration. */
+  async dispose(): Promise<void> {
+    if (this.ownsPool) await this.pool.end();
+  }
+
+  private startInitialization(): Promise<void> {
+    this.initPromise = this.migrationManager.migrate().then(() => {
+      this.initialized = true;
+    });
+    void this.initPromise.catch(() => undefined);
+    return this.initPromise;
+  }
+
+  private async ensureInitialized(): Promise<void> {
+    if (!this.initialized) await this.initialize();
+  }
 
   async create(params: CreateWorkItemRequest): Promise<TaskWorkItem> {
+    await this.ensureInitialized();
     const id = v4();
     const result = await this.pool.query(
       `INSERT INTO lattice_task_work_items
-        (id, tenant_id, task_id, action, actor, thread_id, summary, detail, attempt, workspace_id, project_id)
-       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
+        (id, tenant_id, task_id, action, actor, thread_id, summary, detail, attempt, workspace_id, project_id, created_at)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, date_trunc('milliseconds', clock_timestamp()))
        RETURNING *`,
       [
         id, params.tenantId, params.taskId, params.action, params.actor,
@@ -24,8 +105,36 @@ export class PostgreSQLTaskWorkItemStore implements TaskWorkItemStore {
     return this.rowToItem(result.rows[0]);
   }
 
+  /** Atomically inserts a work item by selecting one exact task snapshot. */
+  async createIfTaskSnapshot(
+    params: CreateWorkItemRequest,
+    snapshot: TaskMutationSnapshot,
+  ): Promise<TaskWorkItem | null> {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `INSERT INTO lattice_task_work_items
+        (id, tenant_id, task_id, action, actor, thread_id, summary, detail, attempt, workspace_id, project_id, created_at)
+       SELECT $1, task.tenant_id, task.id, $2, $3, $4, $5, $6, $7, task.workspace_id, task.project_id,
+         date_trunc('milliseconds', clock_timestamp())
+       FROM lattice_tasks AS task
+       WHERE task.tenant_id = $8 AND task.id = $9 AND task.status = $10
+         AND date_trunc('milliseconds', task.updated_at::timestamptz) = $11::timestamptz
+          AND task.owner_type = $12 AND task.owner_id = $13
+          AND task.workspace_id IS NOT DISTINCT FROM $14 AND task.project_id IS NOT DISTINCT FROM $15
+          AND $14 IS NOT DISTINCT FROM $16 AND $15 IS NOT DISTINCT FROM $17
+       RETURNING *`,
+      [v4(), params.action, params.actor, params.threadId || null, params.summary || null,
+        params.detail ? JSON.stringify(params.detail) : null, params.attempt ?? null,
+        params.tenantId, params.taskId, snapshot.status, new Date(snapshot.updatedAt).toISOString(),
+        snapshot.ownerType, snapshot.ownerId, snapshot.workspaceId, snapshot.projectId,
+        params.workspaceId, params.projectId],
+    );
+    return result.rows[0] ? this.rowToItem(result.rows[0]) : null;
+  }
+
   /** Find an event by its tenant- and task-scoped key without pagination. */
   async findByEventKey(tenantId: string, taskId: string, eventKey: string): Promise<TaskWorkItem | null> {
+    await this.ensureInitialized();
     const result = await this.pool.query(
       `SELECT * FROM lattice_task_work_items
        WHERE tenant_id = $1 AND task_id = $2 AND event_key = $3`,
@@ -36,10 +145,11 @@ export class PostgreSQLTaskWorkItemStore implements TaskWorkItemStore {
 
   /** Atomically return an existing event or create it once. */
   async createIfAbsentByEventKey(params: CreateWorkItemIfAbsentRequest): Promise<TaskWorkItem> {
+    await this.ensureInitialized();
     const result = await this.pool.query(
       `INSERT INTO lattice_task_work_items
-        (id, tenant_id, task_id, action, actor, thread_id, summary, detail, attempt, workspace_id, project_id, event_key)
-       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
+        (id, tenant_id, task_id, action, actor, thread_id, summary, detail, attempt, workspace_id, project_id, event_key, created_at)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, date_trunc('milliseconds', clock_timestamp()))
        ON CONFLICT (tenant_id, task_id, event_key) WHERE event_key IS NOT NULL
        DO UPDATE SET event_key = EXCLUDED.event_key
        RETURNING *`,
@@ -55,6 +165,7 @@ export class PostgreSQLTaskWorkItemStore implements TaskWorkItemStore {
   }
 
   async list(filter: TaskWorkItemListFilter): Promise<TaskWorkItem[]> {
+    await this.ensureInitialized();
     let query = `SELECT * FROM lattice_task_work_items WHERE tenant_id = $1 AND task_id = $2`;
     const params: unknown[] = [filter.tenantId, filter.taskId];
 
@@ -87,6 +198,85 @@ export class PostgreSQLTaskWorkItemStore implements TaskWorkItemStore {
     return result.rows.map((row: Record<string, unknown>) => this.rowToItem(row));
   }
 
+  /** List pending execution results using one bounded PostgreSQL anti-join query. */
+  async listPendingExecutionResults(params: {
+    tenantId: string;
+    taskId: string;
+    limit: number;
+  }): Promise<TaskWorkItem[]> {
+    if (!Number.isSafeInteger(params.limit)
+      || params.limit < 0
+      || params.limit > MAX_PENDING_EXECUTION_RESULTS_LIMIT) {
+      const error = new RangeError(`limit must be a safe integer between 0 and ${MAX_PENDING_EXECUTION_RESULTS_LIMIT}`) as RangeError & {
+        code: "INVALID_LIMIT";
+      };
+      error.code = "INVALID_LIMIT";
+      throw error;
+    }
+    if (params.limit === 0) return [];
+    const result = await this.pool.query(
+      `SELECT result.*
+       FROM lattice_task_work_items AS result
+       WHERE result.tenant_id = $1
+         AND result.task_id = $2
+         AND result.action = 'execution_result'
+         AND result.event_key COLLATE "C" ~ '^execution-result:[A-Za-z0-9._:-]+$'
+         AND NOT EXISTS (
+           SELECT 1
+           FROM lattice_task_work_items AS reconciled
+           WHERE reconciled.tenant_id = result.tenant_id
+             AND reconciled.task_id = result.task_id
+             AND reconciled.action = 'execution_reconciled'
+             AND reconciled.detail ->> 'executionResultId' = result.event_key
+         )
+       ORDER BY result.created_at DESC, result.id DESC
+       LIMIT $3`,
+      [params.tenantId, params.taskId, params.limit],
+    );
+    return result.rows.map((row: Record<string, unknown>) => this.rowToItem(row));
+  }
+
+  /** Lists canonical project lifecycle events with an exclusive cursor. */
+  async listProjectLifecycleEvents(query: ProjectLifecycleEventQuery): Promise<TaskWorkItem[]> {
+    let cursorMilliseconds: number | null = null;
+    if (!Number.isSafeInteger(query.limit) || query.limit < 1 || query.limit > 100
+      || query.actions.length === 0
+      || query.actions.some((action) => !PROJECT_LIFECYCLE_ACTIONS.has(action))) {
+      throw new RangeError("Invalid project lifecycle event page");
+    }
+    if (query.before) {
+      try { cursorMilliseconds = Date.prototype.getTime.call(query.before.createdAt) as number; }
+      catch { throw new RangeError("Invalid project lifecycle event cursor"); }
+      if (!Number.isFinite(cursorMilliseconds) || typeof query.before.id !== "string" || query.before.id.length === 0) {
+        throw new RangeError("Invalid project lifecycle event cursor");
+      }
+    }
+    const cursor = cursorMilliseconds === null ? null : new Date(cursorMilliseconds);
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `SELECT * FROM lattice_task_work_items
+       WHERE tenant_id=$1 AND workspace_id=$2 AND project_id=$3
+         AND action = ANY($4::text[]) AND event_key IS NOT NULL AND event_key <> ''
+         AND ($5::timestamptz IS NULL OR created_at < $5::timestamptz
+           OR (created_at = $5::timestamptz AND id < $6))
+       ORDER BY created_at DESC, id DESC LIMIT $7`,
+      [query.tenantId, query.workspaceId, query.projectId, query.actions,
+        cursor, query.before?.id ?? null, query.limit],
+    );
+    return result.rows.map((row: Record<string, unknown>) => this.assertProjectLifecycleEvent(this.rowToItem(row)));
+  }
+
+  private assertProjectLifecycleEvent(item: TaskWorkItem): TaskWorkItem {
+    let milliseconds: number;
+    try { milliseconds = Date.prototype.getTime.call(item.createdAt) as number; }
+    catch { throw new Error("Invalid project lifecycle event row"); }
+    if (!Number.isFinite(milliseconds) || typeof item.id !== "string" || item.id.length === 0
+      || typeof item.eventKey !== "string" || item.eventKey.length === 0) {
+      throw new Error("Invalid project lifecycle event row");
+    }
+    return { ...item, createdAt: new Date(milliseconds) };
+  }
+
   private rowToItem(row: Record<string, unknown>): TaskWorkItem {
     return {
       id: row.id as string,
@@ -96,7 +286,7 @@ export class PostgreSQLTaskWorkItemStore implements TaskWorkItemStore {
       actor: row.actor as string,
       threadId: row.thread_id as string | undefined,
       summary: row.summary as string | undefined,
-      detail: row.detail as Record<string, unknown> | undefined,
+      detail: isRecord(row.detail) ? row.detail : undefined,
       attempt: row.attempt as number | undefined,
       workspaceId: row.workspace_id as string | undefined,
       projectId: row.project_id as string | undefined,

package/src/stores/ThreadMessageQueueStore.ts

@@ -11,13 +11,21 @@ import type {
   PendingMessage,
   AddMessageParams,
   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 };
 
@@ -58,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) => {
@@ -92,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)
@@ -105,20 +115,61 @@ 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;
+    }
+    const client = await this.pool.connect();
+    try {
+      await client.query("BEGIN");
+      const scope = { tenantId: params.tenantId, assistantId: params.assistantId, workspaceId: params.workspaceId, projectId: params.projectId };
+      const lockKey = `${params.tenantId}:${params.assistantId}:${params.threadId}:${params.workspaceId ?? ""}:${params.projectId ?? ""}`;
+      await client.query("SELECT pg_advisory_xact_lock(hashtext($1))", [lockKey]);
+      const filter = scopeClause(scope, 2);
+      const count = await client.query(
+        `SELECT COUNT(*) as count FROM lattice_thread_message_queue WHERE thread_id = $1 AND status = 'pending'${filter.sql}`,
+        [params.threadId, ...filter.params],
+      );
+      if (parseInt(count.rows[0].count, 10) >= maxSize) {
+        await client.query("ROLLBACK");
+        return false;
+      }
+      const seq = await client.query(
+        `SELECT COALESCE(MAX(sequence_order), 0) + 1 as next_seq FROM lattice_thread_message_queue WHERE thread_id = $1`,
+        [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, 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]);
+    } catch (error) {
+      await client.query("ROLLBACK");
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
   /**
    * Add message at head of queue (high priority, e.g., STEER/Command messages)
    * 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!;
@@ -136,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]);
@@ -148,12 +199,13 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
   /**
    * Get pending messages for thread
    */
-  async getPendingMessages(threadId: string): Promise<PendingMessage[]> {
+  async getPendingMessages(threadId: string, scope?: QueueScope): Promise<PendingMessage[]> {
+    const filter = scopeClause(scope, 2);
     const result = await this.pool.query(
-      `SELECT * FROM lattice_thread_message_queue 
-       WHERE thread_id = $1 AND status = 'pending'
+      `SELECT * FROM lattice_thread_message_queue
+        WHERE thread_id = $1 AND status = 'pending'${filter.sql}
        ORDER BY priority DESC, sequence_order ASC`,
-      [threadId]
+      [threadId, ...filter.params]
     );
 
     return result.rows.map(row => this.rowToMessage(row));
@@ -162,12 +214,13 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
   /**
    * Get processing messages for a thread
    */
-  async getProcessingMessages(threadId: string): Promise<PendingMessage[]> {
+  async getProcessingMessages(threadId: string, scope?: QueueScope): Promise<PendingMessage[]> {
+    const filter = scopeClause(scope, 2);
     const result = await this.pool.query(
       `SELECT * FROM lattice_thread_message_queue 
-       WHERE thread_id = $1 AND status = 'processing'
+        WHERE thread_id = $1 AND status = 'processing'${filter.sql}
        ORDER BY priority DESC, sequence_order ASC`,
-      [threadId]
+      [threadId, ...filter.params]
     );
 
     return result.rows.map(row => this.rowToMessage(row));
@@ -176,11 +229,12 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
   /**
    * Get queue size
    */
-  async getQueueSize(threadId: string): Promise<number> {
+  async getQueueSize(threadId: string, scope?: QueueScope): Promise<number> {
+    const filter = scopeClause(scope, 2);
     const result = await this.pool.query(
       `SELECT COUNT(*) as count FROM lattice_thread_message_queue 
-       WHERE thread_id = $1 AND status = 'pending'`,
-      [threadId]
+        WHERE thread_id = $1 AND status = 'pending'${filter.sql}`,
+      [threadId, ...filter.params]
     );
 
     return parseInt(result.rows[0].count, 10);
@@ -191,28 +245,29 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
    */
   async getThreadsWithPendingMessages(): Promise<ThreadInfo[]> {
     const result = await this.pool.query(
-      `SELECT DISTINCT ON (thread_id) tenant_id, assistant_id, thread_id, workspace_id, project_id
+       `SELECT DISTINCT ON (tenant_id, assistant_id, workspace_id, project_id, thread_id) tenant_id, assistant_id, thread_id, workspace_id, project_id
        FROM lattice_thread_message_queue 
        WHERE status IN ('pending', 'processing')
-       ORDER BY thread_id`
+        ORDER BY tenant_id, assistant_id, workspace_id, project_id, thread_id`
     );
 
     return result.rows.map(row => ({
       tenantId: row.tenant_id,
       assistantId: row.assistant_id,
       threadId: row.thread_id,
-      workspaceId: row.workspace_id || undefined,
-      projectId: row.project_id || undefined,
+      workspaceId: row.workspace_id,
+      projectId: row.project_id,
     }));
   }
 
   /**
    * Remove message
    */
-  async removeMessage(messageId: string): Promise<boolean> {
+  async removeMessage(messageId: string, scope?: QueueScope): Promise<boolean> {
+    const filter = scopeClause(scope, 2);
     const result = await this.pool.query(
-      `DELETE FROM lattice_thread_message_queue WHERE id = $1 RETURNING id`,
-      [messageId]
+      `DELETE FROM lattice_thread_message_queue WHERE id = $1${filter.sql} RETURNING id`,
+      [messageId, ...filter.params]
     );
 
     return (result.rowCount ?? 0) > 0;
@@ -221,20 +276,30 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
   /**
    * Clear all messages for thread
    */
-  async clearMessages(threadId: string): Promise<void> {
+  async clearMessages(threadId: string, scope?: QueueScope): Promise<void> {
+    const filter = scopeClause(scope, 2);
     await this.pool.query(
-      `DELETE FROM lattice_thread_message_queue WHERE thread_id = $1`,
-      [threadId]
+      `DELETE FROM lattice_thread_message_queue WHERE thread_id = $1${filter.sql}`,
+      [threadId, ...filter.params]
     );
   }
 
   /**
    * Mark message as processing
    */
-  async markProcessing(messageId: string): Promise<void> {
+  async markProcessing(messageId: string, customRunConfig?: unknown, scope?: QueueScope): Promise<void> {
+    if (customRunConfig !== undefined) {
+      const filter = scopeClause(scope, 3);
+      await this.pool.query(
+        `UPDATE lattice_thread_message_queue SET status = 'processing', custom_run_config = $2 WHERE id = $1${filter.sql}`,
+        [messageId, JSON.stringify(customRunConfig), ...filter.params]
+      );
+      return;
+    }
+    const filter = scopeClause(scope, 2);
     await this.pool.query(
-      `UPDATE lattice_thread_message_queue SET status = 'processing' WHERE id = $1`,
-      [messageId]
+      `UPDATE lattice_thread_message_queue SET status = 'processing' WHERE id = $1${filter.sql}`,
+      [messageId, ...filter.params]
     );
   }
 
@@ -242,13 +307,14 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
    * Reset all processing messages to pending state for a thread
    * Returns the number of messages reset
    */
-  async resetProcessingToPending(threadId: string): Promise<number> {
+  async resetProcessingToPending(threadId: string, scope?: QueueScope): Promise<number> {
+    const filter = scopeClause(scope, 2);
     const result = await this.pool.query(
       `UPDATE lattice_thread_message_queue 
        SET status = 'pending'
-       WHERE thread_id = $1 AND status = 'processing'
+        WHERE thread_id = $1 AND status = 'processing'${filter.sql}
        RETURNING id`,
-      [threadId]
+      [threadId, ...filter.params]
     );
     return result.rowCount ?? 0;
   }
@@ -267,6 +333,38 @@ 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[] = [];
+  const entries = (["tenantId", "assistantId", "workspaceId", "projectId"] as const)
+    .map((key) => [key, scope[key]] as const);
+  const sql = entries.map(([key, value]) => {
+    const column = key.replace(/[A-Z]/g, (letter) => `_${letter.toLowerCase()}`);
+    if (value == null) return ` AND ${column} IS NULL`;
+    params.push(value);
+    return ` AND ${column} = $${start + params.length - 1}`;
+  }).join("");
+  return { sql, params };
+}