@axiom-lattice/pg-stores 3.3.5 → 3.3.6

package/src/stores/PostgreSQLModelProviderStore.ts

@@ -20,8 +20,12 @@ import { MigrationManager } from '../migrations/migration';
 import {
   createModelProvidersTable,
   createModelProviderModelsTable,
+  addModelProviderApiStyle,
+  addModelProviderModelStale,
+  addModelProviderModelExtra,
+  addModelProviderModelDefault,
 } from '../migrations/model_provider_migrations';
-import { encrypt, decrypt } from '@axiom-lattice/core';
+import { encrypt, safeDecrypt } from '@axiom-lattice/core';
 
 /**
  * PostgreSQL ModelProviderStore options
@@ -76,9 +80,25 @@ interface ModelRow {
   upstream_created: Date | null;
   enabled: boolean;
   stale: boolean;
+  extra: Record<string, unknown> | null;
+  is_default: boolean;
   discovered_at: Date;
 }
 
+/**
+ * A model's stored config, or undefined when it carries none.
+ *
+ * The column is `NOT NULL DEFAULT '{}'`, so "no per-model config" arrives as an
+ * empty object. The protocol types `extra` as optional and the in-memory store
+ * leaves it undefined, so normalising here keeps the two stores
+ * interchangeable rather than letting callers tell which one is behind them.
+ */
+function extraOrUndefined(
+  value: Record<string, unknown> | null
+): Record<string, unknown> | undefined {
+  return value && Object.keys(value).length > 0 ? value : undefined;
+}
+
 /**
  * PostgreSQL implementation of ModelProviderStore
  *
@@ -113,8 +133,17 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
     }
 
     this.migrationManager = new MigrationManager(this.pool);
-    this.migrationManager.register(createModelProvidersTable);
-    this.migrationManager.register(createModelProviderModelsTable);
+    // Every migration that touches these two tables, not just the ones that
+    // create them. This store can migrate itself (standalone `autoMigrate`,
+    // without going through `createPgStoreConfig`), and a list that stops at
+    // v182 leaves such a database without `stale` or `extra` — the reads below
+    // then fail on a missing column rather than returning an empty value.
+    this.migrationManager.register(createModelProvidersTable);     // v181
+    this.migrationManager.register(createModelProviderModelsTable); // v182
+    this.migrationManager.register(addModelProviderApiStyle);       // v183
+    this.migrationManager.register(addModelProviderModelStale);     // v184
+    this.migrationManager.register(addModelProviderModelExtra);     // v191
+    this.migrationManager.register(addModelProviderModelDefault);   // v192
 
     // Auto-migrate by default
     if (options.autoMigrate !== false) {
@@ -440,7 +469,7 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
       ? await this.pool.query<ModelRow>(
           `
           SELECT id, tenant_id, provider_id, model_id, display_name, owned_by,
-                 upstream_created, enabled, stale, discovered_at
+                 upstream_created, enabled, stale, extra, is_default, discovered_at
           FROM lattice_model_provider_models
           WHERE tenant_id = $1 AND provider_id = $2
           ORDER BY model_id ASC
@@ -450,7 +479,7 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
       : await this.pool.query<ModelRow>(
           `
           SELECT id, tenant_id, provider_id, model_id, display_name, owned_by,
-                 upstream_created, enabled, discovered_at
+                 upstream_created, enabled, stale, extra, is_default, discovered_at
           FROM lattice_model_provider_models
           WHERE tenant_id = $1
           ORDER BY model_id ASC
@@ -476,7 +505,7 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
     const result = await this.pool.query<ModelRow>(
       `
       SELECT id, tenant_id, provider_id, model_id, display_name, owned_by,
-             upstream_created, enabled, discovered_at
+             upstream_created, enabled, stale, extra, is_default, discovered_at
       FROM lattice_model_provider_models
       WHERE tenant_id = $1 AND provider_id = $2 AND model_id = $3
     `,
@@ -498,8 +527,12 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
    * 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).
+   * upserted *without* touching `enabled` or `extra`, so the tenant's choices
+   * survive a re-discovery. New rows take the column defaults.
+   *
+   * The default model is the exception to the delete: it is kept and flagged
+   * `stale` instead, because deleting it would take the tenant's default away
+   * with no confirmation.
    */
   async replaceProviderModels(
     tenantId: string,
@@ -513,12 +546,17 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
     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 }>(
+      // Read the current per-row state first, so the write below can state it
+      // explicitly: discovery leaves `enabled` / `extra` out (the tenant's
+      // choice is kept) while an explicit selection supplies it.
+      const current = await client.query<{
+        model_id: string;
+        enabled: boolean;
+        extra: Record<string, unknown> | null;
+        is_default: boolean;
+      }>(
         `
-        SELECT model_id, enabled FROM lattice_model_provider_models
+        SELECT model_id, enabled, extra, is_default FROM lattice_model_provider_models
         WHERE tenant_id = $1 AND provider_id = $2
       `,
         [tenantId, providerId]
@@ -526,6 +564,12 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
       const currentEnabled = new Map(
         current.rows.map((row) => [row.model_id, row.enabled])
       );
+      const currentExtra = new Map(
+        current.rows.map((row) => [row.model_id, row.extra])
+      );
+      const currentIsDefault = new Map(
+        current.rows.map((row) => [row.model_id, row.is_default])
+      );
 
       // De-duplicate by model id: a single `ON CONFLICT DO UPDATE` statement
       // cannot touch the same target row twice, and upstreams (and callers)
@@ -540,7 +584,63 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
         (model) => model.enabled ?? currentEnabled.get(model.id) ?? true
       );
 
-      // Drop only what the upstream no longer offers
+      // Same rule as `enabled`: an omitted `extra` means "leave the stored
+      // config alone", so a sync — which carries no per-model config at all —
+      // cannot wipe what the tenant set by hand.
+      //
+      // The fallback is `{}` rather than null: the column is NOT NULL, and
+      // writing an explicit null fails the constraint for every model that has
+      // no config, which is most of them.
+      const effectiveExtra = rows.map(
+        (model) => model.extra ?? currentExtra.get(model.id) ?? {}
+      );
+
+      // Whether this write is naming a default at all.
+      //
+      // The unique index spans the *tenant*, while this method writes one
+      // provider. Two consequences:
+      //
+      // 1. A swap (A was the default, B now is) would see two true rows
+      //    mid-statement — PostgreSQL checks a unique index per row, so whether
+      //    it fails would depend on which row the batch happened to process
+      //    first. Clearing tenant-wide *before* the upsert makes order
+      //    irrelevant.
+      // 2. Once cleared, an omitted flag must resolve to `false`, not "keep":
+      //    keeping it would restore the old default alongside the new one.
+      const namesDefault = rows.some((model) => model.isDefault === true);
+      if (namesDefault) {
+        await client.query(
+          `UPDATE lattice_model_provider_models SET is_default = FALSE WHERE tenant_id = $1`,
+          [tenantId]
+        );
+      }
+
+      const effectiveIsDefault = rows.map(
+        (model) =>
+          model.isDefault ??
+          (namesDefault ? false : currentIsDefault.get(model.id)) ??
+          false
+      );
+
+      // Drop only what the upstream no longer offers — except the default.
+      //
+      // Deleting it would take the tenant's default away with no confirmation,
+      // and nothing would put it back. That is not hypothetical: a discovery
+      // pass replaces the set wholesale, so an upstream that stops offering the
+      // default model would silently delete it. Keep it, flagged, the way the
+      // tenant-facing sync does.
+
+      await client.query(
+        `
+        UPDATE lattice_model_provider_models
+           SET stale = TRUE
+         WHERE tenant_id = $1
+           AND provider_id = $2
+           AND model_id <> ALL($3::text[])
+           AND is_default
+      `,
+        [tenantId, providerId, rows.map((model) => model.id)]
+      );
 
       await client.query(
         `
@@ -548,26 +648,28 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
         WHERE tenant_id = $1
           AND provider_id = $2
           AND model_id <> ALL($3::text[])
+          AND NOT is_default
       `,
         [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.
+        // `enabled` and `extra` are computed above so that an omitted value
+        // resolves to the row's current one rather than being cleared.
         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
+            upstream_created, raw, enabled, stale, extra, is_default, 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[]
+            $6::text[], $7::timestamp[], $8::jsonb[], $9::boolean[], $10::boolean[],
+            $11::jsonb[], $12::boolean[], $13::timestamp[]
           ) AS t(
             id, tenant_id, provider_id, model_id, display_name, owned_by,
-            upstream_created, raw, enabled, stale, discovered_at
+            upstream_created, raw, enabled, stale, extra, is_default, discovered_at
           )
           ON CONFLICT (tenant_id, provider_id, model_id) DO UPDATE SET
             display_name = EXCLUDED.display_name,
@@ -576,6 +678,8 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
             raw = EXCLUDED.raw,
             enabled = EXCLUDED.enabled,
             stale = EXCLUDED.stale,
+            extra = EXCLUDED.extra,
+            is_default = EXCLUDED.is_default,
             discovered_at = EXCLUDED.discovered_at
         `,
           [
@@ -593,6 +697,8 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
             ),
             effectiveEnabled,
             rows.map((model) => model.stale ?? false),
+            effectiveExtra.map((extra) => JSON.stringify(extra)),
+            effectiveIsDefault,
             rows.map(() => new Date().toISOString()),
           ]
         );
@@ -682,12 +788,7 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
     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');
-      }
+      apiKey = safeDecrypt(row.api_key_enc);
     }
 
     return {
@@ -725,6 +826,8 @@ export class PostgreSQLModelProviderStore implements ModelProviderStore {
       upstreamCreatedAt: row.upstream_created || undefined,
       enabled: row.enabled,
       stale: row.stale,
+      extra: extraOrUndefined(row.extra),
+      isDefault: row.is_default,
       discoveredAt: row.discovered_at,
     };
   }