@axiom-lattice/pg-stores 3.0.0 → 3.0.2

package/src/stores/PostgreSQLAgentWebAppStore.ts

@@ -0,0 +1,280 @@
+import { randomUUID } from "crypto";
+import type {
+  AgentWebApp,
+  AgentWebAppAppearance,
+  AgentWebAppFeatures,
+  AgentWebAppScope,
+  AgentWebAppStatus,
+  AgentWebAppStore,
+  AgentWebAppStorePatch,
+  AgentWebAppUpdateOptions,
+  CreateAgentWebAppInput,
+} from "@axiom-lattice/protocols";
+import { Pool } from "pg";
+import type { PoolConfig } from "pg";
+
+import { createAgentWebAppsTable } from "../migrations/agent_web_apps_migration";
+import { MigrationManager } from "../migrations/migration";
+
+export interface PostgreSQLAgentWebAppStoreOptions {
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
+  autoMigrate?: boolean;
+}
+
+type AgentWebAppRow = {
+  id: string;
+  tenant_id: string;
+  assistant_id: string;
+  name: string;
+  description: string | null;
+  status: AgentWebAppStatus;
+  integration: unknown;
+  scope: unknown;
+  features: unknown;
+  appearance: unknown;
+  created_at: Date;
+  updated_at: Date;
+};
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function isStringArray(value: unknown): value is string[] {
+  return Array.isArray(value) && value.every((item) => typeof item === "string");
+}
+
+function isIntegration(value: unknown): value is AgentWebApp["integration"] {
+  return isRecord(value) && value.type === "react_sdk";
+}
+
+function isScope(value: unknown): value is AgentWebAppScope {
+  return isRecord(value)
+    && typeof value.defaultProjectId === "string"
+    && isStringArray(value.allowedProjectIds)
+    && (value.defaultModelKey === undefined || typeof value.defaultModelKey === "string")
+    && (value.allowedModelKeys === undefined || isStringArray(value.allowedModelKeys));
+}
+
+function isFeatures(value: unknown): value is AgentWebAppFeatures {
+  return isRecord(value)
+    && typeof value.projectSelector === "boolean"
+    && typeof value.modelSelector === "boolean"
+    && typeof value.threadManagement === "boolean"
+    && typeof value.attachments === "boolean"
+    && typeof value.hitl === "boolean"
+    && typeof value.genUI === "boolean";
+}
+
+function isAppearance(value: unknown): value is AgentWebAppAppearance {
+  return isRecord(value)
+    && (value.title === undefined || typeof value.title === "string")
+    && (value.welcomeMessage === undefined || typeof value.welcomeMessage === "string")
+    && (value.primaryColor === undefined || typeof value.primaryColor === "string");
+}
+
+function isStatus(value: unknown): value is AgentWebAppStatus {
+  return value === "draft" || value === "active" || value === "disabled";
+}
+
+function mapRow(row: AgentWebAppRow): AgentWebApp {
+  if (
+    !isStatus(row.status)
+    || !isIntegration(row.integration)
+    || !isScope(row.scope)
+    || !isFeatures(row.features)
+    || !isAppearance(row.appearance)
+  ) {
+    throw new Error(`Invalid agent web app row: ${row.id}`);
+  }
+
+  return {
+    id: row.id,
+    tenantId: row.tenant_id,
+    assistantId: row.assistant_id,
+    name: row.name,
+    description: row.description ?? undefined,
+    status: row.status,
+    integration: row.integration,
+    scope: row.scope,
+    features: row.features,
+    appearance: row.appearance,
+    createdAt: row.created_at,
+    updatedAt: row.updated_at,
+  };
+}
+
+export class PostgreSQLAgentWebAppStore implements AgentWebAppStore {
+  private pool: Pool;
+  private migrationManager!: MigrationManager;
+  private initialized = false;
+  private ownsPool = true;
+  private initPromise: Promise<void> | null = null;
+
+  constructor(options: PostgreSQLAgentWebAppStoreOptions) {
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      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.migrationManager = new MigrationManager(this.pool);
+    this.migrationManager.register(createAgentWebAppsTable);
+
+    if (options.autoMigrate !== false) {
+      this.startInitialization();
+    }
+  }
+
+  async initialize(): Promise<void> {
+    if (this.initialized) return;
+    if (this.initPromise) return this.initPromise;
+    return this.startInitialization();
+  }
+
+  private startInitialization(): Promise<void> {
+    this.initPromise = this.migrationManager.migrate().then(() => {
+      this.initialized = true;
+    });
+    // Eager migration errors remain on initPromise for initialize/operations.
+    void this.initPromise.catch(() => undefined);
+    return this.initPromise;
+  }
+
+  async dispose(): Promise<void> {
+    if (this.ownsPool) await this.pool.end();
+  }
+
+  private async ensureInitialized(): Promise<void> {
+    if (!this.initialized) await this.initialize();
+  }
+
+  async list(tenantId: string, assistantId?: string): Promise<AgentWebApp[]> {
+    await this.ensureInitialized();
+    const result = assistantId === undefined
+      ? await this.pool.query<AgentWebAppRow>(
+        `SELECT * FROM lattice_agent_web_apps
+         WHERE tenant_id = $1
+         ORDER BY created_at DESC, id DESC`,
+        [tenantId],
+      )
+      : await this.pool.query<AgentWebAppRow>(
+        `SELECT * FROM lattice_agent_web_apps
+         WHERE tenant_id = $1 AND assistant_id = $2
+         ORDER BY created_at DESC, id DESC`,
+        [tenantId, assistantId],
+      );
+    return result.rows.map(mapRow);
+  }
+
+  async getById(tenantId: string, webAppId: string): Promise<AgentWebApp | null> {
+    await this.ensureInitialized();
+    const result = await this.pool.query<AgentWebAppRow>(
+      `SELECT * FROM lattice_agent_web_apps
+       WHERE tenant_id = $1 AND id = $2`,
+      [tenantId, webAppId],
+    );
+    return result.rows[0] ? mapRow(result.rows[0]) : null;
+  }
+
+  async findById(webAppId: string): Promise<AgentWebApp | null> {
+    await this.ensureInitialized();
+    const result = await this.pool.query<AgentWebAppRow>(
+      `SELECT * FROM lattice_agent_web_apps
+       WHERE id = $1`,
+      [webAppId],
+    );
+    return result.rows[0] ? mapRow(result.rows[0]) : null;
+  }
+
+  async create(tenantId: string, input: CreateAgentWebAppInput): Promise<AgentWebApp> {
+    await this.ensureInitialized();
+    const id = `webapp_${randomUUID().replace(/-/g, "")}`;
+    const result = await this.pool.query<AgentWebAppRow>(
+      `INSERT INTO lattice_agent_web_apps
+       (id, tenant_id, assistant_id, name, description, status, integration, scope, features, appearance)
+       VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8::jsonb, $9::jsonb, $10::jsonb)
+       RETURNING *`,
+      [
+        id,
+        tenantId,
+        input.assistantId,
+        input.name,
+        input.description ?? null,
+        "draft",
+        JSON.stringify(input.integration),
+        JSON.stringify(input.scope),
+        JSON.stringify(input.features),
+        JSON.stringify(input.appearance),
+      ],
+    );
+    return mapRow(result.rows[0]);
+  }
+
+  async update(
+    tenantId: string,
+    webAppId: string,
+    patch: AgentWebAppStorePatch,
+    options?: AgentWebAppUpdateOptions,
+  ): Promise<AgentWebApp | null> {
+    await this.ensureInitialized();
+    const assignments: string[] = [];
+    const values: unknown[] = [];
+    const add = (column: string, value: unknown): void => {
+      values.push(value);
+      assignments.push(`${column} = $${values.length}`);
+    };
+
+    if (patch.name !== undefined) add("name", patch.name);
+    if (Object.prototype.hasOwnProperty.call(patch, "description")) {
+      add("description", patch.description ?? null);
+    }
+    if (patch.scope !== undefined) {
+      add("scope", JSON.stringify(patch.scope));
+    }
+    if (patch.features !== undefined) {
+      add("features", JSON.stringify(patch.features));
+    }
+    if (patch.appearance !== undefined) {
+      add("appearance", JSON.stringify(patch.appearance));
+    }
+    if (patch.status !== undefined) add("status", patch.status);
+
+    if (assignments.length === 0) return this.getById(tenantId, webAppId);
+    values.push(tenantId, webAppId);
+    const tenantParam = values.length - 1;
+    const idParam = values.length;
+    const expectedUpdatedAtClause = options?.expectedUpdatedAt
+      ? ` AND date_trunc('milliseconds', updated_at) = $${values.push(options.expectedUpdatedAt)}`
+      : "";
+    const result = await this.pool.query<AgentWebAppRow>(
+      `UPDATE lattice_agent_web_apps
+       SET ${assignments.join(", ")}, updated_at = GREATEST(
+         date_trunc('milliseconds', clock_timestamp()),
+         updated_at + interval '1 millisecond'
+       )
+       WHERE tenant_id = $${tenantParam} AND id = $${idParam}${expectedUpdatedAtClause}
+       RETURNING *`,
+      values,
+    );
+    return result.rows[0] ? mapRow(result.rows[0]) : null;
+  }
+
+  async delete(tenantId: string, webAppId: string): Promise<boolean> {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `DELETE FROM lattice_agent_web_apps
+       WHERE tenant_id = $1 AND id = $2`,
+      [tenantId, webAppId],
+    );
+    return result.rowCount === 1;
+  }
+}