@axiom-lattice/pg-stores 1.0.70 → 1.0.72

package/src/stores/PostgreSQLChannelInstallationStore.ts

@@ -132,6 +132,32 @@ export class PostgreSQLChannelInstallationStore
     return result.rows.map((row) => this.mapRowToInstallation(row));
   }
 
+  async getAllInstallations(
+    channel?: ChannelInstallationType,
+  ): Promise<ChannelInstallation[]> {
+    await this.ensureInitialized();
+
+    const result = channel
+      ? await this.pool.query<ChannelInstallationRow>(
+          `
+          SELECT id, tenant_id, channel, name, config, enabled, fallback_agent_id, reject_when_no_binding, created_at, updated_at
+          FROM lattice_channel_installations
+          WHERE channel = $1
+          ORDER BY created_at DESC
+          `,
+          [channel],
+        )
+      : await this.pool.query<ChannelInstallationRow>(
+          `
+          SELECT id, tenant_id, channel, name, config, enabled, fallback_agent_id, reject_when_no_binding, created_at, updated_at
+          FROM lattice_channel_installations
+          ORDER BY created_at DESC
+          `,
+        );
+
+    return result.rows.map((row) => this.mapRowToInstallation(row));
+  }
+
   async createInstallation(
     tenantId: string,
     installationId: string,

package/src/stores/PostgreSQLTaskStore.ts

@@ -0,0 +1,294 @@
+/**
+ * PostgreSQL implementation of TaskStore
+ */
+
+import { Pool } from "pg";
+import type { PoolConfig } from "pg";
+import type {
+  TaskStore,
+  TaskItem,
+  CreateTaskRequest,
+  UpdateTaskRequest,
+  TaskListFilter,
+} from "@axiom-lattice/protocols";
+import { MigrationManager } from "../migrations/migration";
+import { createTasksTable } from "../migrations/task_migration";
+import { v4 as uuidv4 } from "uuid";
+
+interface TaskRow {
+  id: string;
+  tenant_id: string;
+  owner_type: string;
+  owner_id: string;
+  title: string;
+  description: string | null;
+  status: string;
+  priority: string;
+  due_date: string | null;
+  metadata: Record<string, unknown> | null;
+  parent_id: string | null;
+  source_id: string | null;
+  context: Record<string, unknown> | null;
+  created_at: string;
+  updated_at: string;
+}
+
+function mapRowToTask(row: TaskRow): TaskItem {
+  return {
+    id: row.id,
+    tenantId: row.tenant_id,
+    ownerType: row.owner_type as "user" | "agent",
+    ownerId: row.owner_id,
+    title: row.title,
+    description: row.description ?? undefined,
+    status: row.status as TaskItem["status"],
+    priority: row.priority as TaskItem["priority"],
+    dueDate: row.due_date ?? undefined,
+    metadata: row.metadata ?? undefined,
+    parentId: row.parent_id ?? undefined,
+    sourceId: row.source_id ?? undefined,
+    context: row.context ?? undefined,
+    createdAt: new Date(row.created_at),
+    updatedAt: new Date(row.updated_at),
+  };
+}
+
+export interface PostgreSQLTaskStoreOptions {
+  poolConfig: string | PoolConfig;
+  autoMigrate?: boolean;
+}
+
+export class PostgreSQLTaskStore implements TaskStore {
+  private pool: Pool;
+  private migrationManager: MigrationManager;
+  private initialized: boolean = false;
+  private ownsPool: boolean = true;
+  private initPromise: Promise<void> | null = null;
+
+  constructor(options: PostgreSQLTaskStoreOptions) {
+    if (typeof options.poolConfig === "string") {
+      this.pool = new Pool({ connectionString: options.poolConfig });
+    } else {
+      this.pool = new Pool(options.poolConfig);
+    }
+
+    this.migrationManager = new MigrationManager(this.pool);
+    this.migrationManager.register(createTasksTable);
+
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error("Failed to initialize PostgreSQLTaskStore:", error);
+        throw error;
+      });
+    }
+  }
+
+  async initialize(): Promise<void> {
+    if (this.initialized) {
+      return;
+    }
+
+    if (this.initPromise) {
+      return this.initPromise;
+    }
+
+    this.initPromise = (async () => {
+      try {
+        await this.migrationManager.migrate();
+        this.initialized = true;
+      } finally {
+        this.initPromise = null;
+      }
+    })();
+
+    return this.initPromise;
+  }
+
+  async create(
+    params: CreateTaskRequest & { tenantId: string; ownerType: string; ownerId: string },
+  ): Promise<TaskItem> {
+    await this.ensureInitialized();
+    const id = uuidv4();
+    const now = new Date().toISOString();
+
+    await this.pool.query(
+      `INSERT INTO lattice_tasks (id, tenant_id, owner_type, owner_id, title, description, status, priority, due_date, metadata, parent_id, source_id, context, created_at, updated_at)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)`,
+      [
+        id,
+        params.tenantId,
+        params.ownerType,
+        params.ownerId,
+        params.title,
+        params.description || null,
+        params.status || "pending",
+        params.priority || "medium",
+        params.dueDate || null,
+        params.metadata ? JSON.stringify(params.metadata) : null,
+        params.parentId || null,
+        params.sourceId || null,
+        params.context ? JSON.stringify(params.context) : null,
+        now,
+        now,
+      ],
+    );
+
+    return (await this.getById(params.tenantId, id))!;
+  }
+
+  async getById(tenantId: string, id: string): Promise<TaskItem | null> {
+    await this.ensureInitialized();
+    const result = await this.pool.query<TaskRow>(
+      `SELECT * FROM lattice_tasks WHERE tenant_id = $1 AND id = $2`,
+      [tenantId, id],
+    );
+    if (result.rows.length === 0) return null;
+    return mapRowToTask(result.rows[0]);
+  }
+
+  async list(filter: TaskListFilter): Promise<TaskItem[]> {
+    await this.ensureInitialized();
+    const conditions: string[] = [`tenant_id = $1`];
+    const params: unknown[] = [filter.tenantId];
+    let paramIndex = 2;
+
+    if (filter.ownerType) {
+      conditions.push(`owner_type = $${paramIndex++}`);
+      params.push(filter.ownerType);
+    }
+    if (filter.ownerId) {
+      conditions.push(`owner_id = $${paramIndex++}`);
+      params.push(filter.ownerId);
+    }
+    if (filter.status) {
+      conditions.push(`status = $${paramIndex++}`);
+      params.push(filter.status);
+    }
+    if (filter.priority) {
+      conditions.push(`priority = $${paramIndex++}`);
+      params.push(filter.priority);
+    }
+    if (filter.parentId) {
+      conditions.push(`parent_id = $${paramIndex++}`);
+      params.push(filter.parentId);
+    }
+    if (filter.sourceId) {
+      conditions.push(`source_id = $${paramIndex++}`);
+      params.push(filter.sourceId);
+    }
+    if (filter.metadata) {
+      for (const [key, value] of Object.entries(filter.metadata)) {
+        const safeKey = key.replace(/[^a-zA-Z0-9_]/g, '');
+        if (safeKey) {
+          conditions.push(`metadata->>'${safeKey}' = $${paramIndex++}`);
+          params.push(String(value));
+        }
+      }
+    }
+
+    const where = conditions.join(" AND ");
+    const limit = filter.limit || 100;
+    const offset = filter.offset || 0;
+
+    const result = await this.pool.query<TaskRow>(
+      `SELECT * FROM lattice_tasks WHERE ${where} ORDER BY created_at DESC LIMIT $${paramIndex++} OFFSET $${paramIndex++}`,
+      [...params, limit, offset],
+    );
+    return result.rows.map((r) => mapRowToTask(r));
+  }
+
+  async update(
+    tenantId: string,
+    id: string,
+    updates: UpdateTaskRequest,
+  ): Promise<TaskItem | null> {
+    await this.ensureInitialized();
+    const existing = await this.getById(tenantId, id);
+    if (!existing) return null;
+
+    const setClauses: string[] = [];
+    const params: unknown[] = [];
+    let paramIndex = 1;
+
+    if (updates.title !== undefined) {
+      setClauses.push(`title = $${paramIndex++}`);
+      params.push(updates.title);
+    }
+    if (updates.description !== undefined) {
+      setClauses.push(`description = $${paramIndex++}`);
+      params.push(updates.description);
+    }
+    if (updates.status !== undefined) {
+      setClauses.push(`status = $${paramIndex++}`);
+      params.push(updates.status);
+    }
+    if (updates.priority !== undefined) {
+      setClauses.push(`priority = $${paramIndex++}`);
+      params.push(updates.priority);
+    }
+    if (updates.dueDate !== undefined) {
+      setClauses.push(`due_date = $${paramIndex++}`);
+      params.push(updates.dueDate);
+    }
+    if (updates.metadata !== undefined) {
+      setClauses.push(`metadata = $${paramIndex++}`);
+      params.push(JSON.stringify(updates.metadata));
+    }
+    if (updates.parentId !== undefined) {
+      setClauses.push(`parent_id = $${paramIndex++}`);
+      params.push(updates.parentId);
+    }
+    if (updates.sourceId !== undefined) {
+      setClauses.push(`source_id = $${paramIndex++}`);
+      params.push(updates.sourceId);
+    }
+    if (updates.context !== undefined) {
+      setClauses.push(`context = $${paramIndex++}`);
+      params.push(JSON.stringify(updates.context));
+    }
+    if (updates.ownerType !== undefined) {
+      setClauses.push(`owner_type = $${paramIndex++}`);
+      params.push(updates.ownerType);
+    }
+    if (updates.ownerId !== undefined) {
+      setClauses.push(`owner_id = $${paramIndex++}`);
+      params.push(updates.ownerId);
+    }
+
+    if (setClauses.length === 0) {
+      return existing;
+    }
+
+    setClauses.push(`updated_at = $${paramIndex++}`);
+    params.push(new Date().toISOString());
+
+    params.push(tenantId, id);
+    await this.pool.query(
+      `UPDATE lattice_tasks SET ${setClauses.join(", ")} WHERE tenant_id = $${paramIndex++} AND id = $${paramIndex++}`,
+      params,
+    );
+
+    return this.getById(tenantId, id);
+  }
+
+  async delete(tenantId: string, id: string): Promise<boolean> {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `DELETE FROM lattice_tasks WHERE tenant_id = $1 AND id = $2`,
+      [tenantId, id],
+    );
+    return (result.rowCount ?? 0) > 0;
+  }
+
+  async dispose(): Promise<void> {
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
+  }
+
+  private async ensureInitialized(): Promise<void> {
+    if (!this.initialized) {
+      await this.initialize();
+    }
+  }
+}

package/src/stores/PostgreSQLWorkflowTrackingStore.ts

@@ -16,7 +16,7 @@ import {
   WorkflowRunStatus,
 } from "@axiom-lattice/protocols";
 import { MigrationManager } from "../migrations/migration";
-import { createWorkflowTrackingTables } from "../migrations/workflow_tracking_migrations";
+import { createWorkflowTrackingTables, addStepThreadId } from "../migrations/workflow_tracking_migrations";
 
 export interface PostgreSQLWorkflowTrackingStoreOptions {
   poolConfig: string | PoolConfig;
@@ -38,6 +38,7 @@ export class PostgreSQLWorkflowTrackingStore implements WorkflowTrackingStore {
 
     this.migrationManager = new MigrationManager(this.pool);
     this.migrationManager.register(createWorkflowTrackingTables);
+    this.migrationManager.register(addStepThreadId);
 
     if (options.autoMigrate !== false) {
       this.initialize().catch((error) => {
@@ -171,23 +172,47 @@ export class PostgreSQLWorkflowTrackingStore implements WorkflowTrackingStore {
     const id = `step_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`;
 
     await this.pool.query(
-      `INSERT INTO lattice_workflow_steps (id, run_id, tenant_id, step_type, step_name, edge_from, edge_to, edge_purpose, input, status, started_at, created_at, updated_at)
-       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)`,
+      `INSERT INTO lattice_workflow_steps (id, run_id, tenant_id, step_type, step_name, edge_from, edge_to, edge_purpose, input, thread_id, status, started_at, created_at, updated_at)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)`,
       [id, request.runId, request.tenantId, request.stepType, request.stepName,
        request.edgeFrom || null, request.edgeTo || null, request.edgePurpose || null,
        request.input ? JSON.stringify(request.input) : null,
+       request.threadId || null,
        'running', now, now, now]
     );
 
     return {
       id, runId: request.runId, tenantId: request.tenantId,
       stepType: request.stepType, stepName: request.stepName,
+      threadId: request.threadId,
       edgeFrom: request.edgeFrom, edgeTo: request.edgeTo, edgePurpose: request.edgePurpose,
       input: request.input, status: 'running', startedAt: now,
       createdAt: now, updatedAt: now,
     };
   }
 
+  async upsertRunStep(request: CreateRunStepRequest): Promise<RunStep> {
+    await this.ensureInitialized();
+    // Look up by (run_id, step_type, step_name) — the idempotency key
+    const result = await this.pool.query(
+      `SELECT * FROM lattice_workflow_steps WHERE run_id = $1 AND step_type = $2 AND step_name = $3`,
+      [request.runId, request.stepType, request.stepName]
+    );
+    if (result.rows.length > 0) {
+      const existing = this.mapRowToRunStep(result.rows[0]);
+      // Update threadId if not already set and request has one
+      if (!existing.threadId && request.threadId) {
+        await this.pool.query(
+          `UPDATE lattice_workflow_steps SET thread_id = $1, updated_at = $2 WHERE run_id = $3 AND id = $4`,
+          [request.threadId, new Date(), request.runId, existing.id]
+        );
+        existing.threadId = request.threadId;
+      }
+      return existing;
+    }
+    return this.createRunStep(request);
+  }
+
   async updateRunStep(runId: string, stepId: string, updates: UpdateRunStepRequest): Promise<RunStep | null> {
     await this.ensureInitialized();
     const now = new Date();
@@ -281,6 +306,7 @@ export class PostgreSQLWorkflowTrackingStore implements WorkflowTrackingStore {
       tenantId: row.tenant_id,
       stepType: row.step_type as StepType,
       stepName: row.step_name,
+      threadId: row.thread_id,
       edgeFrom: row.edge_from,
       edgeTo: row.edge_to,
       edgePurpose: row.edge_purpose,

package/src/stores/ThreadMessageQueueStore.ts

@@ -17,6 +17,7 @@ import { createThreadMessageQueueTable } from "../migrations/thread_message_queu
 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";
 
 export type { PendingMessage, AddMessageParams, ThreadInfo };
 
@@ -45,6 +46,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
     this.migrationManager.register(addPriorityAndCommandColumns);
     this.migrationManager.register(addCustomRunConfigColumn);
     this.migrationManager.register(alterMessageQueueIdColumn);
+    this.migrationManager.register(addWorkspaceProjectToQueue);
 
     if (options.autoMigrate !== false) {
       this.initialize().catch((error) => {
@@ -79,7 +81,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
    * Add message to queue
    */
   async addMessage(params: AddMessageParams): Promise<PendingMessage> {
-    const { threadId, tenantId, assistantId, content, type = "human", priority = 0, command, custom_run_config, id } = 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)
     const seqResult = await this.pool.query(
@@ -92,10 +94,10 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
 
     const result = await this.pool.query(
       `INSERT INTO lattice_thread_message_queue 
-       (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)
+       (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 *`,
-      [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]
+      [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]
     );
 
     return this.rowToMessage(result.rows[0]);
@@ -106,7 +108,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
    * Uses priority=100 to ensure message is processed first
    */
   async addMessageAtHead(params: AddMessageParams): Promise<PendingMessage> {
-    const { threadId, tenantId, assistantId, content, type = "human", command, custom_run_config, id } = params;
+    const { threadId, tenantId, assistantId, workspaceId, projectId, content, type = "human", command, custom_run_config, id } = params;
 
     const resolvedTenantId = tenantId!;
     const resolvedAssistantId = assistantId!;
@@ -123,10 +125,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, message_content, message_type, sequence_order, priority, command, custom_run_config)
-       VALUES ($1, $2, $3, $4, $5, $6, $7, 100, $8, $9)
+       (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)
        RETURNING *`,
-      [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]
+      [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]
     );
 
     return this.rowToMessage(result.rows[0]);
@@ -178,7 +180,7 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
    */
   async getThreadsWithPendingMessages(): Promise<ThreadInfo[]> {
     const result = await this.pool.query(
-      `SELECT DISTINCT tenant_id, assistant_id, thread_id 
+      `SELECT DISTINCT ON (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`
@@ -188,6 +190,8 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
       tenantId: row.tenant_id,
       assistantId: row.assistant_id,
       threadId: row.thread_id,
+      workspaceId: row.workspace_id || undefined,
+      projectId: row.project_id || undefined,
     }));
   }