@axiom-lattice/pg-stores 3.1.21 → 3.2.1

package/src/stores/PostgreSQLModelProviderStore.ts

@@ -0,0 +1,731 @@
+/**
+ * PostgreSQL implementation of ModelProviderStore
+ */
+
+import { Pool, PoolClient } from 'pg';
+import type { PoolConfig } from 'pg';
+import {
+  ModelProviderStore,
+  ModelProviderEntry,
+  ModelProviderModelEntry,
+  ModelProviderStatus,
+  ModelProviderWireProtocol,
+  ModelProviderApiStyle,
+  ModelProviderLlmProvider,
+  CreateModelProviderRequest,
+  UpdateModelProviderRequest,
+  ProviderModelInput,
+} from '@axiom-lattice/protocols';
+import { MigrationManager } from '../migrations/migration';
+import {
+  createModelProvidersTable,
+  createModelProviderModelsTable,
+} from '../migrations/model_provider_migrations';
+import { encrypt, decrypt } from '@axiom-lattice/core';
+
+/**
+ * PostgreSQL ModelProviderStore options
+ */
+export interface PostgreSQLModelProviderStoreOptions {
+  /**
+   * External pool instance. When provided, poolConfig is not required and the
+   * caller is responsible for pool lifecycle and migrations.
+   */
+  pool?: Pool;
+
+  /**
+   * PostgreSQL connection pool configuration
+   * Can be a connection string or PoolConfig object
+   */
+  poolConfig?: string | PoolConfig;
+
+  /**
+   * Whether to run migrations automatically on initialization
+   * @default true
+   */
+  autoMigrate?: boolean;
+}
+
+interface ProviderRow {
+  id: string;
+  tenant_id: string;
+  name: string;
+  display_name: string | null;
+  protocol: string;
+  api_style: string;
+  llm_provider: string;
+  base_url: string;
+  api_key_enc: string | null;
+  api_key_hint: string | null;
+  enabled: boolean;
+  status: string;
+  last_error: string | null;
+  last_discovered_at: Date | null;
+  created_by: string | null;
+  created_at: Date;
+  updated_at: Date;
+}
+
+interface ModelRow {
+  id: string;
+  tenant_id: string;
+  provider_id: string;
+  model_id: string;
+  display_name: string | null;
+  owned_by: string | null;
+  upstream_created: Date | null;
+  enabled: boolean;
+  stale: boolean;
+  discovered_at: Date;
+}
+
+/**
+ * PostgreSQL implementation of ModelProviderStore
+ *
+ * Features:
+ * - Multi-tenant isolation via tenant_id
+ * - Automatic API key encryption/decryption
+ * - Unique constraint on (tenant_id, name) and (tenant_id, model_id)
+ */
+export class PostgreSQLModelProviderStore implements ModelProviderStore {
+  private pool: Pool;
+  private migrationManager!: MigrationManager;
+  private initialized: boolean = false;
+  private ownsPool: boolean = true;
+  private initPromise: Promise<void> | null = null;
+
+  constructor(options: PostgreSQLModelProviderStoreOptions) {
+    // Use externally-managed pool if provided
+    if (options.pool) {
+      this.pool = options.pool;
+      this.ownsPool = false;
+      this.initialized = true;
+      return;
+    }
+
+    // Create Pool from config
+    if (typeof options.poolConfig === 'string') {
+      this.pool = new Pool({ connectionString: options.poolConfig });
+    } else if (options.poolConfig) {
+      this.pool = new Pool(options.poolConfig as PoolConfig);
+    } else {
+      throw new Error('Either pool or poolConfig must be provided');
+    }
+
+    this.migrationManager = new MigrationManager(this.pool);
+    this.migrationManager.register(createModelProvidersTable);
+    this.migrationManager.register(createModelProviderModelsTable);
+
+    // Auto-migrate by default
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error('Failed to initialize PostgreSQLModelProviderStore:', error);
+        throw error;
+      });
+    }
+  }
+
+  /**
+   * Initialize the store and run migrations
+   * Uses a promise-based lock to prevent concurrent initialization
+   */
+  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;
+  }
+
+  /**
+   * Get all provider configurations for a tenant
+   */
+  async listProviders(tenantId: string): Promise<ModelProviderEntry[]> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<ProviderRow>(
+      `
+      SELECT id, tenant_id, name, display_name, protocol, api_style, llm_provider, base_url,
+             api_key_enc, api_key_hint, enabled, status, last_error,
+             last_discovered_at, created_by, created_at, updated_at
+      FROM lattice_model_providers
+      WHERE tenant_id = $1
+      ORDER BY created_at DESC
+    `,
+      [tenantId]
+    );
+
+    return result.rows.map((row) => this.mapRowToEntry(row));
+  }
+
+  /**
+   * Get a provider configuration by ID
+   */
+  async getProviderById(
+    tenantId: string,
+    id: string
+  ): Promise<ModelProviderEntry | null> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<ProviderRow>(
+      `
+      SELECT id, tenant_id, name, display_name, protocol, api_style, llm_provider, base_url,
+             api_key_enc, api_key_hint, enabled, status, last_error,
+             last_discovered_at, created_by, created_at, updated_at
+      FROM lattice_model_providers
+      WHERE tenant_id = $1 AND id = $2
+    `,
+      [tenantId, id]
+    );
+
+    if (result.rows.length === 0) {
+      return null;
+    }
+
+    return this.mapRowToEntry(result.rows[0]);
+  }
+
+  /**
+   * Get a provider configuration by business name
+   */
+  async getProviderByName(
+    tenantId: string,
+    name: string
+  ): Promise<ModelProviderEntry | null> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<ProviderRow>(
+      `
+      SELECT id, tenant_id, name, display_name, protocol, api_style, llm_provider, base_url,
+             api_key_enc, api_key_hint, enabled, status, last_error,
+             last_discovered_at, created_by, created_at, updated_at
+      FROM lattice_model_providers
+      WHERE tenant_id = $1 AND name = $2
+    `,
+      [tenantId, name]
+    );
+
+    if (result.rows.length === 0) {
+      return null;
+    }
+
+    return this.mapRowToEntry(result.rows[0]);
+  }
+
+  /**
+   * Create a new provider configuration
+   */
+  async createProvider(
+    tenantId: string,
+    id: string,
+    data: CreateModelProviderRequest
+  ): Promise<ModelProviderEntry> {
+    await this.ensureInitialized();
+
+    const now = new Date();
+    const nowString = now.toISOString();
+    const apiKeyEnc = data.apiKey ? encrypt(data.apiKey) : null;
+    const apiKeyHint = data.apiKey ? data.apiKey.slice(-4) : null;
+
+    await this.pool.query(
+      `
+      INSERT INTO lattice_model_providers (
+        id, tenant_id, name, display_name, protocol, api_style, llm_provider, base_url,
+        api_key_enc, api_key_hint, enabled, status, created_at, updated_at
+      )
+      VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, 'unknown', $12::timestamp, $13::timestamp)
+      ON CONFLICT (tenant_id, id) DO UPDATE SET
+        name = EXCLUDED.name,
+        display_name = EXCLUDED.display_name,
+        protocol = EXCLUDED.protocol,
+        api_style = EXCLUDED.api_style,
+        llm_provider = EXCLUDED.llm_provider,
+        base_url = EXCLUDED.base_url,
+        api_key_enc = EXCLUDED.api_key_enc,
+        api_key_hint = EXCLUDED.api_key_hint,
+        enabled = EXCLUDED.enabled,
+        updated_at = EXCLUDED.updated_at
+    `,
+      [
+        id,
+        tenantId,
+        data.name,
+        data.displayName || null,
+        data.protocol,
+        data.apiStyle || 'chat-completions',
+        // Phase 1 pins the LLM provider: the UI does not offer a choice, and
+        // some initChatModel branches hardcode their own baseURL and would
+        // silently ignore the tenant's (see ADR-138 decision 7b).
+        'openai',
+        data.baseURL,
+        apiKeyEnc,
+        apiKeyHint,
+        data.enabled !== false,
+        nowString,
+        nowString,
+      ]
+    );
+
+    const created = await this.getProviderById(tenantId, id);
+    if (!created) {
+      throw new Error('Failed to create model provider');
+    }
+
+    return created;
+  }
+
+  /**
+   * Update an existing provider configuration.
+   * An omitted or empty `apiKey` leaves the stored key untouched.
+   */
+  async updateProvider(
+    tenantId: string,
+    id: string,
+    updates: Partial<UpdateModelProviderRequest>
+  ): Promise<ModelProviderEntry | null> {
+    await this.ensureInitialized();
+
+    const existing = await this.getProviderById(tenantId, id);
+    if (!existing) {
+      return null;
+    }
+
+    const updateData: Record<string, unknown> = {};
+
+    if (updates.name !== undefined) {
+      updateData.name = updates.name;
+    }
+
+    if (updates.displayName !== undefined) {
+      updateData.display_name = updates.displayName || null;
+    }
+
+    if (updates.protocol !== undefined) {
+      updateData.protocol = updates.protocol;
+    }
+
+    if (updates.apiStyle !== undefined) {
+      updateData.api_style = updates.apiStyle;
+    }
+
+    if (updates.baseURL !== undefined) {
+      updateData.base_url = updates.baseURL;
+    }
+
+    // An empty or omitted apiKey means "keep the existing key"; `clearApiKey`
+    // is the only way to remove one outright.
+    if (updates.clearApiKey === true) {
+      updateData.api_key_enc = null;
+      updateData.api_key_hint = null;
+    } else if (updates.apiKey) {
+      updateData.api_key_enc = encrypt(updates.apiKey);
+      updateData.api_key_hint = updates.apiKey.slice(-4);
+    }
+
+    if (updates.enabled !== undefined) {
+      updateData.enabled = updates.enabled;
+    }
+
+    if (Object.keys(updateData).length === 0) {
+      return existing;
+    }
+
+    // Always update the updated_at timestamp
+    updateData.updated_at = new Date().toISOString();
+
+    const fields = Object.keys(updateData);
+    const values: unknown[] = Object.values(updateData);
+
+    // Add WHERE clause values at the end
+    values.push(tenantId);
+    values.push(id);
+
+    const setClauses = fields.map((field, index) =>
+      field === 'updated_at'
+        ? `${field} = $${index + 1}::timestamp`
+        : `${field} = $${index + 1}`
+    );
+
+    const whereTenantIndex = fields.length + 1;
+    const whereIdIndex = fields.length + 2;
+
+    const sql = `
+      UPDATE lattice_model_providers
+      SET ${setClauses.join(', ')}
+      WHERE tenant_id = $${whereTenantIndex} AND id = $${whereIdIndex}
+    `;
+
+    await this.pool.query(sql, values);
+
+    return await this.getProviderById(tenantId, id);
+  }
+
+  /**
+   * Delete a provider configuration and its model snapshot
+   */
+  async deleteProvider(tenantId: string, id: string): Promise<boolean> {
+    await this.ensureInitialized();
+
+    await this.deleteModelsByProvider(tenantId, id);
+
+    const result = await this.pool.query(
+      `
+      DELETE FROM lattice_model_providers
+      WHERE tenant_id = $1 AND id = $2
+    `,
+      [tenantId, id]
+    );
+
+    return result.rowCount !== null && result.rowCount > 0;
+  }
+
+  /**
+   * Record the outcome of a discovery attempt
+   */
+  async markProviderStatus(
+    tenantId: string,
+    id: string,
+    status: ModelProviderStatus,
+    lastError?: string,
+    discoveredAt?: Date
+  ): Promise<void> {
+    await this.ensureInitialized();
+
+    await this.pool.query(
+      `
+      UPDATE lattice_model_providers
+      SET status = $3,
+          last_error = $4,
+          last_discovered_at = CASE
+            WHEN $5::timestamp IS NULL THEN last_discovered_at
+            ELSE $5::timestamp
+          END,
+          updated_at = $6::timestamp
+      WHERE tenant_id = $1 AND id = $2
+    `,
+      [
+        tenantId,
+        id,
+        status,
+        status === 'ok' ? null : lastError || null,
+        discoveredAt ? discoveredAt.toISOString() : null,
+        new Date().toISOString(),
+      ]
+    );
+  }
+
+  /**
+   * List snapshot models for a tenant
+   */
+  async listModels(
+    tenantId: string,
+    providerId?: string
+  ): Promise<ModelProviderModelEntry[]> {
+    await this.ensureInitialized();
+
+    const result = providerId
+      ? await this.pool.query<ModelRow>(
+          `
+          SELECT id, tenant_id, provider_id, model_id, display_name, owned_by,
+                 upstream_created, enabled, stale, discovered_at
+          FROM lattice_model_provider_models
+          WHERE tenant_id = $1 AND provider_id = $2
+          ORDER BY model_id ASC
+        `,
+          [tenantId, providerId]
+        )
+      : await this.pool.query<ModelRow>(
+          `
+          SELECT id, tenant_id, provider_id, model_id, display_name, owned_by,
+                 upstream_created, enabled, discovered_at
+          FROM lattice_model_provider_models
+          WHERE tenant_id = $1
+          ORDER BY model_id ASC
+        `,
+          [tenantId]
+        );
+
+    return result.rows.map((row) => this.mapRowToModel(row));
+  }
+
+  /**
+   * Get a single snapshot model.
+   * Keyed by provider as well as model id, since two providers in the same
+   * tenant may expose the same upstream model id.
+   */
+  async getModel(
+    tenantId: string,
+    providerId: string,
+    modelId: string
+  ): Promise<ModelProviderModelEntry | null> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<ModelRow>(
+      `
+      SELECT id, tenant_id, provider_id, model_id, display_name, owned_by,
+             upstream_created, enabled, discovered_at
+      FROM lattice_model_provider_models
+      WHERE tenant_id = $1 AND provider_id = $2 AND model_id = $3
+    `,
+      [tenantId, providerId, modelId]
+    );
+
+    if (result.rows.length === 0) {
+      return null;
+    }
+
+    return this.mapRowToModel(result.rows[0]);
+  }
+
+  /**
+   * Replace a provider's model snapshot with the result of one discovery.
+   *
+   * Runs in a single transaction. Model ids are scoped per provider, so the
+   * same upstream id may appear under several providers in one tenant — each
+   * is an independent model.
+   *
+   * Models the upstream no longer offers are deleted; models that survive are
+   * upserted *without* touching `enabled`, so the tenant's enable/disable
+   * choice survives a re-discovery. New rows take the column default (enabled).
+   */
+  async replaceProviderModels(
+    tenantId: string,
+    providerId: string,
+    models: ProviderModelInput[]
+  ): Promise<string[]> {
+    await this.ensureInitialized();
+
+    const client = await this.pool.connect();
+
+    try {
+      await client.query('BEGIN');
+
+      // Read the current enabled state first, so the write below can state it
+      // explicitly: discovery leaves `enabled` out (the tenant's choice is kept)
+      // while an explicit selection supplies it.
+      const current = await client.query<{ model_id: string; enabled: boolean }>(
+        `
+        SELECT model_id, enabled FROM lattice_model_provider_models
+        WHERE tenant_id = $1 AND provider_id = $2
+      `,
+        [tenantId, providerId]
+      );
+      const currentEnabled = new Map(
+        current.rows.map((row) => [row.model_id, row.enabled])
+      );
+
+      // De-duplicate by model id: a single `ON CONFLICT DO UPDATE` statement
+      // cannot touch the same target row twice, and upstreams (and callers)
+      // do repeat ids. Last occurrence wins.
+      const deduped = new Map<string, ProviderModelInput>();
+      for (const model of models) {
+        deduped.set(model.id, model);
+      }
+      const rows = [...deduped.values()];
+
+      const effectiveEnabled = rows.map(
+        (model) => model.enabled ?? currentEnabled.get(model.id) ?? true
+      );
+
+      // Drop only what the upstream no longer offers
+
+      await client.query(
+        `
+        DELETE FROM lattice_model_provider_models
+        WHERE tenant_id = $1
+          AND provider_id = $2
+          AND model_id <> ALL($3::text[])
+      `,
+        [tenantId, providerId, rows.map((model) => model.id)]
+      );
+
+      if (models.length > 0) {
+        // One statement for the whole snapshot rather than one per model.
+        // `enabled` is deliberately absent from both the column list and the
+        // update clause, so existing rows keep their value.
+        await client.query(
+          `
+          INSERT INTO lattice_model_provider_models (
+            id, tenant_id, provider_id, model_id, display_name, owned_by,
+            upstream_created, raw, enabled, stale, discovered_at
+          )
+          SELECT * FROM unnest(
+            $1::text[], $2::text[], $3::text[], $4::text[], $5::text[],
+            $6::text[], $7::timestamp[], $8::jsonb[], $9::boolean[], $10::boolean[], $11::timestamp[]
+          ) AS t(
+            id, tenant_id, provider_id, model_id, display_name, owned_by,
+            upstream_created, raw, enabled, stale, discovered_at
+          )
+          ON CONFLICT (tenant_id, provider_id, model_id) DO UPDATE SET
+            display_name = EXCLUDED.display_name,
+            owned_by = EXCLUDED.owned_by,
+            upstream_created = EXCLUDED.upstream_created,
+            raw = EXCLUDED.raw,
+            enabled = EXCLUDED.enabled,
+            stale = EXCLUDED.stale,
+            discovered_at = EXCLUDED.discovered_at
+        `,
+          [
+            rows.map((model) => `${providerId}:${model.id}`),
+            rows.map(() => tenantId),
+            rows.map(() => providerId),
+            rows.map((model) => model.id),
+            rows.map((model) => model.displayName ?? null),
+            rows.map((model) => model.ownedBy ?? null),
+            rows.map((model) =>
+              model.createdAt ? new Date(model.createdAt).toISOString() : null
+            ),
+            rows.map((model) =>
+              model.raw === undefined ? null : JSON.stringify(model.raw)
+            ),
+            effectiveEnabled,
+            rows.map((model) => model.stale ?? false),
+            rows.map(() => new Date().toISOString()),
+          ]
+        );
+      }
+
+      await client.query('COMMIT');
+    } catch (error) {
+      await client.query('ROLLBACK');
+      throw error;
+    } finally {
+      client.release();
+    }
+
+    return models.map((model) => model.id);
+  }
+
+  /**
+   * Enable or disable a single snapshot model
+   */
+  async setModelEnabled(
+    tenantId: string,
+    providerId: string,
+    modelId: string,
+    enabled: boolean
+  ): Promise<ModelProviderModelEntry | null> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query(
+      `
+      UPDATE lattice_model_provider_models
+      SET enabled = $4
+      WHERE tenant_id = $1 AND provider_id = $2 AND model_id = $3
+    `,
+      [tenantId, providerId, modelId, enabled]
+    );
+
+    if (result.rowCount === null || result.rowCount === 0) {
+      return null;
+    }
+
+    return await this.getModel(tenantId, providerId, modelId);
+  }
+
+  /**
+   * Delete every snapshot model belonging to a provider
+   */
+  async deleteModelsByProvider(
+    tenantId: string,
+    providerId: string
+  ): Promise<number> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query(
+      `
+      DELETE FROM lattice_model_provider_models
+      WHERE tenant_id = $1 AND provider_id = $2
+    `,
+      [tenantId, providerId]
+    );
+
+    return result.rowCount ?? 0;
+  }
+
+  /**
+   * Dispose resources and close the connection pool
+   */
+  async dispose(): Promise<void> {
+    if (this.ownsPool && this.pool) {
+      await this.pool.end();
+    }
+  }
+
+  /**
+   * Ensure store is initialized
+   */
+  private async ensureInitialized(): Promise<void> {
+    if (!this.initialized) {
+      await this.initialize();
+    }
+  }
+
+  /**
+   * Map provider row to ModelProviderEntry
+   * Automatically decrypts the API key if present
+   */
+  private mapRowToEntry(row: ProviderRow): ModelProviderEntry {
+    let apiKey: string | undefined;
+
+    if (row.api_key_enc) {
+      try {
+        apiKey = decrypt(row.api_key_enc);
+      } catch (error) {
+        console.error('Failed to decrypt model provider API key:', error);
+        throw new Error('Failed to decrypt model provider API key');
+      }
+    }
+
+    return {
+      id: row.id,
+      tenantId: row.tenant_id,
+      name: row.name,
+      displayName: row.display_name || undefined,
+      protocol: row.protocol as ModelProviderWireProtocol,
+      apiStyle: row.api_style as ModelProviderApiStyle,
+      llmProvider: row.llm_provider as ModelProviderLlmProvider,
+      baseURL: row.base_url,
+      apiKey,
+      apiKeyHint: row.api_key_hint || undefined,
+      enabled: row.enabled,
+      status: row.status as ModelProviderStatus,
+      lastError: row.last_error || undefined,
+      lastDiscoveredAt: row.last_discovered_at || undefined,
+      createdBy: row.created_by || undefined,
+      createdAt: row.created_at,
+      updatedAt: row.updated_at,
+    };
+  }
+
+  /**
+   * Map model row to ModelProviderModelEntry
+   */
+  private mapRowToModel(row: ModelRow): ModelProviderModelEntry {
+    return {
+      id: row.id,
+      tenantId: row.tenant_id,
+      providerId: row.provider_id,
+      modelId: row.model_id,
+      displayName: row.display_name || undefined,
+      ownedBy: row.owned_by || undefined,
+      upstreamCreatedAt: row.upstream_created || undefined,
+      enabled: row.enabled,
+      stale: row.stale,
+      discoveredAt: row.discovered_at,
+    };
+  }
+}