@axiom-lattice/pg-stores 1.0.47 → 1.0.49

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "1.0.47",
+  "version": "1.0.49",
   "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.30",
-    "@axiom-lattice/core": "2.1.57"
+    "@axiom-lattice/protocols": "2.1.31",
+    "@axiom-lattice/core": "2.1.59"
   },
   "devDependencies": {
     "@types/node": "^20.11.24",

package/src/index.ts

@@ -42,9 +42,11 @@ export * from "./migrations/tenant_migrations";
 export * from "./migrations/thread_message_queue_migrations";
 export * from "./migrations/channel_identity_mapping_migration";
 export * from "./migrations/channel_installation_migrations";
+export * from "./migrations/workflow_tracking_migrations";
 export * from "./stores/ThreadMessageQueueStore";
 export * from "./stores/ChannelIdentityMappingStore";
 export * from "./stores/PostgreSQLChannelInstallationStore";
+export * from "./stores/PostgreSQLWorkflowTrackingStore";
 
 // Re-export for convenience
 export { PostgreSQLThreadStore } from "./stores/PostgreSQLThreadStore";
@@ -61,6 +63,7 @@ export { PostgreSQLTenantStore } from "./stores/PostgreSQLTenantStore";
 export { PostgreSQLUserTenantLinkStore } from "./stores/PostgreSQLUserTenantLinkStore";
 export { ChannelIdentityMappingStore } from "./stores/ChannelIdentityMappingStore";
 export { PostgreSQLChannelInstallationStore } from "./stores/PostgreSQLChannelInstallationStore";
+export { PostgreSQLWorkflowTrackingStore } from "./stores/PostgreSQLWorkflowTrackingStore";
 
 // Re-export types from protocols
 export type {
@@ -120,4 +123,11 @@ export type {
   UserTenantRole,
   CreateUserTenantLinkRequest,
   UpdateUserTenantLinkRequest,
+  WorkflowTrackingStore,
+  WorkflowRun,
+  RunStep,
+  CreateWorkflowRunRequest,
+  UpdateWorkflowRunRequest,
+  CreateRunStepRequest,
+  UpdateRunStepRequest,
 } from "@axiom-lattice/protocols";

package/src/migrations/workflow_tracking_migrations.ts

@@ -0,0 +1,79 @@
+/**
+ * PostgreSQL migrations for workflow tracking tables
+ */
+
+import { Migration } from "./migration";
+
+export const createWorkflowTrackingTables: Migration = {
+  name: "create_workflow_tracking_tables",
+  version: 103,
+
+  up: async (client) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_workflow_runs (
+        id VARCHAR(255) PRIMARY KEY,
+        tenant_id VARCHAR(255) NOT NULL,
+        assistant_id VARCHAR(255) NOT NULL,
+        thread_id VARCHAR(255) NOT NULL,
+        status VARCHAR(50) NOT NULL DEFAULT 'running',
+        topology_edges JSONB NOT NULL DEFAULT '[]',
+        total_edges INTEGER NOT NULL DEFAULT 0,
+        completed_edges INTEGER NOT NULL DEFAULT 0,
+        error_message TEXT,
+        metadata JSONB DEFAULT '{}',
+        started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+        completed_at TIMESTAMPTZ,
+        created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+        updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+      );
+
+      CREATE TABLE IF NOT EXISTS lattice_workflow_steps (
+        id VARCHAR(255) NOT NULL,
+        run_id VARCHAR(255) NOT NULL REFERENCES lattice_workflow_runs(id) ON DELETE CASCADE,
+        tenant_id VARCHAR(255) NOT NULL,
+        step_type VARCHAR(50) NOT NULL,
+        step_name VARCHAR(255) NOT NULL,
+        edge_from VARCHAR(255),
+        edge_to VARCHAR(255),
+        edge_purpose TEXT,
+        input JSONB,
+        output JSONB,
+        status VARCHAR(50) NOT NULL DEFAULT 'running',
+        error_message TEXT,
+        started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+        completed_at TIMESTAMPTZ,
+        duration_ms INTEGER,
+        created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+        updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+        PRIMARY KEY (id, run_id)
+      );
+
+      CREATE INDEX IF NOT EXISTS idx_workflow_runs_tenant
+        ON lattice_workflow_runs (tenant_id);
+      CREATE INDEX IF NOT EXISTS idx_workflow_runs_thread
+        ON lattice_workflow_runs (tenant_id, thread_id);
+      CREATE INDEX IF NOT EXISTS idx_workflow_runs_assistant
+        ON lattice_workflow_runs (tenant_id, assistant_id);
+      CREATE INDEX IF NOT EXISTS idx_workflow_runs_status
+        ON lattice_workflow_runs (status);
+
+      CREATE INDEX IF NOT EXISTS idx_workflow_steps_run
+        ON lattice_workflow_steps (run_id);
+      CREATE INDEX IF NOT EXISTS idx_workflow_steps_type_status
+        ON lattice_workflow_steps (run_id, step_type, status);
+    `);
+  },
+
+  down: async (client) => {
+    await client.query(`
+      DROP INDEX IF EXISTS idx_workflow_steps_type_status;
+      DROP INDEX IF EXISTS idx_workflow_steps_run;
+      DROP INDEX IF EXISTS idx_workflow_runs_status;
+      DROP INDEX IF EXISTS idx_workflow_runs_assistant;
+      DROP INDEX IF EXISTS idx_workflow_runs_thread;
+      DROP INDEX IF EXISTS idx_workflow_runs_tenant;
+      DROP TABLE IF EXISTS lattice_workflow_steps;
+      DROP TABLE IF EXISTS lattice_workflow_runs;
+    `);
+  },
+};

package/src/stores/PostgreSQLWorkflowTrackingStore.ts

@@ -0,0 +1,293 @@
+/**
+ * PostgreSQL implementation of WorkflowTrackingStore
+ */
+
+import { Pool } from "pg";
+import type { PoolConfig } from "pg";
+import {
+  WorkflowTrackingStore,
+  WorkflowRun,
+  RunStep,
+  CreateWorkflowRunRequest,
+  UpdateWorkflowRunRequest,
+  CreateRunStepRequest,
+  UpdateRunStepRequest,
+  StepType,
+  WorkflowRunStatus,
+} from "@axiom-lattice/protocols";
+import { MigrationManager } from "../migrations/migration";
+import { createWorkflowTrackingTables } from "../migrations/workflow_tracking_migrations";
+
+export interface PostgreSQLWorkflowTrackingStoreOptions {
+  poolConfig: string | PoolConfig;
+  autoMigrate?: boolean;
+}
+
+export class PostgreSQLWorkflowTrackingStore implements WorkflowTrackingStore {
+  private pool: Pool;
+  private migrationManager: MigrationManager;
+  private initialized: boolean = false;
+  private initPromise: Promise<void> | null = null;
+
+  constructor(options: PostgreSQLWorkflowTrackingStoreOptions) {
+    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(createWorkflowTrackingTables);
+
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error("Failed to initialize PostgreSQLWorkflowTrackingStore:", error);
+        throw error;
+      });
+    }
+  }
+
+  async dispose(): Promise<void> {
+    await this.pool.end();
+  }
+
+  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;
+  }
+
+  private async ensureInitialized(): Promise<void> {
+    if (!this.initialized) await this.initialize();
+  }
+
+  async createWorkflowRun(request: CreateWorkflowRunRequest): Promise<WorkflowRun> {
+    await this.ensureInitialized();
+    const now = new Date();
+    const id = `${request.threadId}_${Date.now()}`;
+
+    await this.pool.query(
+      `INSERT INTO lattice_workflow_runs (id, tenant_id, assistant_id, thread_id, status, topology_edges, total_edges, completed_edges, metadata, started_at, created_at, updated_at)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)`,
+      [id, request.tenantId, request.assistantId, request.threadId, 'running',
+       JSON.stringify(request.topologyEdges), request.topologyEdges.length, 0,
+       JSON.stringify(request.metadata || {}), now, now, now]
+    );
+
+    return {
+      id, status: 'running', startedAt: now, completedAt: undefined,
+      ...request, completedEdges: 0, totalEdges: request.topologyEdges.length,
+      createdAt: now, updatedAt: now,
+    };
+  }
+
+  async getWorkflowRun(runId: string): Promise<WorkflowRun | null> {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `SELECT * FROM lattice_workflow_runs WHERE id = $1`, [runId]
+    );
+    if (result.rows.length === 0) return null;
+    return this.mapRowToWorkflowRun(result.rows[0]);
+  }
+
+  async updateWorkflowRun(runId: string, updates: UpdateWorkflowRunRequest): Promise<WorkflowRun | null> {
+    await this.ensureInitialized();
+    const now = new Date();
+    const setClauses: string[] = ['updated_at = $2'];
+    const values: any[] = [runId, now];
+    let paramIdx = 3;
+
+    if (updates.status !== undefined) {
+      setClauses.push(`status = $${paramIdx++}`);
+      values.push(updates.status);
+    }
+    if (updates.completedEdges !== undefined) {
+      setClauses.push(`completed_edges = $${paramIdx++}`);
+      values.push(updates.completedEdges);
+    }
+    if (updates.errorMessage !== undefined) {
+      setClauses.push(`error_message = $${paramIdx++}`);
+      values.push(updates.errorMessage);
+    }
+    if (updates.completedAt !== undefined) {
+      setClauses.push(`completed_at = $${paramIdx++}`);
+      values.push(updates.completedAt);
+    }
+    if (updates.metadata !== undefined) {
+      setClauses.push(`metadata = $${paramIdx++}`);
+      values.push(JSON.stringify(updates.metadata));
+    }
+
+    const query = `UPDATE lattice_workflow_runs SET ${setClauses.join(', ')} WHERE id = $1`;
+    await this.pool.query(query, values);
+    return this.getWorkflowRun(runId);
+  }
+
+  async getWorkflowRunsByThreadId(tenantId: string, threadId: string): Promise<WorkflowRun[]> {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `SELECT * FROM lattice_workflow_runs WHERE tenant_id = $1 AND thread_id = $2 ORDER BY created_at DESC`,
+      [tenantId, threadId]
+    );
+    return result.rows.map(this.mapRowToWorkflowRun);
+  }
+
+  async getWorkflowRunsByAssistantId(tenantId: string, assistantId: string): Promise<WorkflowRun[]> {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `SELECT * FROM lattice_workflow_runs WHERE tenant_id = $1 AND assistant_id = $2 ORDER BY created_at DESC`,
+      [tenantId, assistantId]
+    );
+    return result.rows.map(this.mapRowToWorkflowRun);
+  }
+
+  async getWorkflowRunsByTenantId(tenantId: string): Promise<WorkflowRun[]> {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `SELECT * FROM lattice_workflow_runs WHERE tenant_id = $1 ORDER BY created_at DESC`,
+      [tenantId]
+    );
+    return result.rows.map(this.mapRowToWorkflowRun);
+  }
+
+  async createRunStep(request: CreateRunStepRequest): Promise<RunStep> {
+    await this.ensureInitialized();
+    const now = new Date();
+    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)`,
+      [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,
+       'running', now, now, now]
+    );
+
+    return {
+      id, runId: request.runId, tenantId: request.tenantId,
+      stepType: request.stepType, stepName: request.stepName,
+      edgeFrom: request.edgeFrom, edgeTo: request.edgeTo, edgePurpose: request.edgePurpose,
+      input: request.input, status: 'running', startedAt: now,
+      createdAt: now, updatedAt: now,
+    };
+  }
+
+  async updateRunStep(runId: string, stepId: string, updates: UpdateRunStepRequest): Promise<RunStep | null> {
+    await this.ensureInitialized();
+    const now = new Date();
+    const setClauses: string[] = ['updated_at = $3'];
+    const values: any[] = [runId, stepId, now];
+    let paramIdx = 4;
+
+    if (updates.status !== undefined) {
+      setClauses.push(`status = $${paramIdx++}`);
+      values.push(updates.status);
+    }
+    if (updates.output !== undefined) {
+      setClauses.push(`output = $${paramIdx++}`);
+      values.push(JSON.stringify(updates.output));
+    }
+    if (updates.errorMessage !== undefined) {
+      setClauses.push(`error_message = $${paramIdx++}`);
+      values.push(updates.errorMessage);
+    }
+    if (updates.completedAt !== undefined) {
+      setClauses.push(`completed_at = $${paramIdx++}`);
+      values.push(updates.completedAt);
+    }
+    if (updates.durationMs !== undefined) {
+      setClauses.push(`duration_ms = $${paramIdx++}`);
+      values.push(updates.durationMs);
+    }
+
+    const query = `UPDATE lattice_workflow_steps SET ${setClauses.join(', ')} WHERE run_id = $1 AND id = $2`;
+    const result = await this.pool.query(query, values);
+
+    if (result.rowCount === 0) return null;
+
+    const row = await this.pool.query(
+      `SELECT * FROM lattice_workflow_steps WHERE run_id = $1 AND id = $2`,
+      [runId, stepId]
+    );
+    return row.rows.length > 0 ? this.mapRowToRunStep(row.rows[0]) : null;
+  }
+
+  async getRunSteps(runId: string): Promise<RunStep[]> {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `SELECT * FROM lattice_workflow_steps WHERE run_id = $1 ORDER BY created_at ASC`,
+      [runId]
+    );
+    return result.rows.map(this.mapRowToRunStep);
+  }
+
+  async getRunStepsByType(runId: string, stepType: StepType): Promise<RunStep[]> {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `SELECT * FROM lattice_workflow_steps WHERE run_id = $1 AND step_type = $2 ORDER BY created_at ASC`,
+      [runId, stepType]
+    );
+    return result.rows.map(this.mapRowToRunStep);
+  }
+
+  async getInterruptedSteps(runId: string): Promise<RunStep[]> {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `SELECT * FROM lattice_workflow_steps WHERE run_id = $1 AND status = 'interrupted' ORDER BY created_at ASC`,
+      [runId]
+    );
+    return result.rows.map(this.mapRowToRunStep);
+  }
+
+  private mapRowToWorkflowRun(row: any): WorkflowRun {
+    return {
+      id: row.id,
+      tenantId: row.tenant_id,
+      assistantId: row.assistant_id,
+      threadId: row.thread_id,
+      status: row.status as WorkflowRunStatus,
+      topologyEdges: typeof row.topology_edges === 'string' ? JSON.parse(row.topology_edges) : row.topology_edges || [],
+      totalEdges: row.total_edges,
+      completedEdges: row.completed_edges,
+      errorMessage: row.error_message,
+      metadata: typeof row.metadata === 'string' ? JSON.parse(row.metadata) : row.metadata || {},
+      startedAt: row.started_at,
+      completedAt: row.completed_at,
+      createdAt: row.created_at,
+      updatedAt: row.updated_at,
+    };
+  }
+
+  private mapRowToRunStep(row: any): RunStep {
+    return {
+      id: row.id,
+      runId: row.run_id,
+      tenantId: row.tenant_id,
+      stepType: row.step_type as StepType,
+      stepName: row.step_name,
+      edgeFrom: row.edge_from,
+      edgeTo: row.edge_to,
+      edgePurpose: row.edge_purpose,
+      input: typeof row.input === 'string' ? JSON.parse(row.input) : row.input,
+      output: typeof row.output === 'string' ? JSON.parse(row.output) : row.output,
+      status: row.status,
+      errorMessage: row.error_message,
+      startedAt: row.started_at,
+      completedAt: row.completed_at,
+      durationMs: row.duration_ms,
+      createdAt: row.created_at,
+      updatedAt: row.updated_at,
+    };
+  }
+}