@axiom-lattice/pg-stores 5.1.2 → 5.2.0

package/src/__tests__/PostgreSQLUserTenantLinkStore.test.ts

@@ -0,0 +1,62 @@
+import { randomUUID } from "node:crypto";
+import { Pool } from "pg";
+import {
+  PLATFORM_TENANT_ID,
+  PLATFORM_TENANT_PROTECTED_ERROR,
+} from "@axiom-lattice/protocols";
+import { createUserTenantLinkTable } from "../migrations/user_tenant_link_migrations";
+import { PostgreSQLUserTenantLinkStore } from "../stores/PostgreSQLUserTenantLinkStore";
+
+const databaseUrl = process.env.TEST_DATABASE_URL;
+const describeIfDatabase = databaseUrl ? describe : describe.skip;
+
+describeIfDatabase("PostgreSQLUserTenantLinkStore platform tenant guard", () => {
+  let pool: Pool;
+  let store: PostgreSQLUserTenantLinkStore;
+  const prefix = `utl-${randomUUID()}`;
+
+  beforeAll(async () => {
+    pool = new Pool({ connectionString: databaseUrl });
+    const client = await pool.connect();
+    try {
+      await createUserTenantLinkTable.up(client);
+    } finally {
+      client.release();
+    }
+    store = new PostgreSQLUserTenantLinkStore({ pool });
+  });
+
+  afterEach(async () => {
+    await pool.query(
+      `DELETE FROM lattice_user_tenant_links WHERE user_id LIKE $1`,
+      [`${prefix}%`],
+    );
+  });
+
+  afterAll(async () => {
+    await pool.query(
+      `DELETE FROM lattice_user_tenant_links WHERE user_id LIKE $1`,
+      [`${prefix}%`],
+    );
+    await pool.end();
+  });
+
+  it("still creates links for ordinary tenants", async () => {
+    await expect(
+      store.createLink({ userId: `${prefix}-u1`, tenantId: "acme" }),
+    ).resolves.toMatchObject({ tenantId: "acme", role: "member" });
+  });
+
+  it("rejects creating a link to the hidden platform tenant", async () => {
+    await expect(
+      store.createLink({
+        userId: `${prefix}-u2`,
+        tenantId: PLATFORM_TENANT_ID,
+      }),
+    ).rejects.toMatchObject({ code: PLATFORM_TENANT_PROTECTED_ERROR });
+
+    await expect(
+      store.getLink(`${prefix}-u2`, PLATFORM_TENANT_ID),
+    ).resolves.toBeNull();
+  });
+});

package/src/createPgStoreConfig.ts

@@ -63,7 +63,7 @@ import { createChannelBindingsTable } from "./migrations/channel_bindings_migrat
 import { createDatabaseConfigsTable } from "./migrations/database_config_migrations";
 import { createConnectionConfigsTable, addConnectionSecretFields, reencryptLegacyConnectionSecrets } from "./migrations/connection_config_migrations";
 import { createMetricsConfigsTable } from "./migrations/metrics_config_migrations";
-import { createMcpServerConfigsTable } from "./migrations/mcp_server_config_migrations";
+import { createMcpServerConfigsTable, createPersonalMcpServerConfigsTable } from "./migrations/mcp_server_config_migrations";
 import { registerThreadMessageQueueMigrations } from "./migrations/queue_migrations";
 import { createWorkflowTrackingTables, addStepThreadId, addWorkflowRunsTenantStatusUpdatedIndex } from "./migrations/workflow_tracking_migrations";
 import { evalMigrations } from "./migrations/eval_migrations";
@@ -83,6 +83,7 @@ import { createInviteLinksTable } from "./migrations/invite_link_migration";
 import { addOidcIdentityClaimInviteFields } from "./migrations/oidc_identity_claim_invite_migration";
 import { createPlatformSettingsTable } from "./migrations/platform_settings_migration";
 import { addOidcProviderAllowedHosts } from "./migrations/oidc_provider_allowed_hosts_migration";
+import { createPlatformAuditTable } from "./migrations/platform_audit_migration";
 import { taskMigrations } from "./migrations/task_migration";
 import { createMenuItemsTable } from "./migrations/menu_items_migration";
 import { addFileContentType } from "./migrations/menu_items_add_file_type";
@@ -141,6 +142,7 @@ import { PostgreSQLOutboxStore } from "./stores/PostgreSQLOutboxStore";
 import { PostgreSQLPluginDeploymentStore } from "./stores/PostgreSQLPluginDeploymentStore";
 import { PostgreSQLAgentRunLedger } from "./stores/PostgreSQLAgentRunLedger";
 import { PostgreSQLPlatformSettingsStore } from "./stores/PostgreSQLPlatformSettingsStore";
+import { PostgreSQLPlatformAuditStore } from "./stores/PostgreSQLPlatformAuditStore";
 import { PostgreSQLChannelReplyJobStore } from "./stores/PostgreSQLChannelReplyJobStore";
 import { PostgreSQLCheckpointStorage } from "./stores/PostgreSQLCheckpointStorage";
 import { PostgreSQLStore } from "./stores/PostgreSQLStore";
@@ -253,6 +255,8 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(addOidcIdentityClaimInviteFields); // v216
   mm.register(createPlatformSettingsTable);     // v217
   mm.register(addOidcProviderAllowedHosts);      // v218
+  mm.register(createPlatformAuditTable);         // v219
+  mm.register(createPersonalMcpServerConfigsTable); // v220
 
   await mm.migrate();
 
@@ -281,6 +285,7 @@ export async function createPgStoreConfig(connectionString: string) {
     oidcIdentityClaim: new PostgreSQLOidcIdentityClaimStore(opts),
     inviteLink: new PostgreSQLInviteLinkStore(opts),
     platformSettings: new PostgreSQLPlatformSettingsStore(opts),
+    platformAudit: new PostgreSQLPlatformAuditStore(opts),
     channelBinding: new ChannelBindingStore(opts),
     channelInstallation: new PostgreSQLChannelInstallationStore(opts),
     thread: new PostgreSQLThreadStore(opts),

package/src/index.ts

@@ -40,6 +40,8 @@ export * from "./migrations/oidc_identity_claim_invite_migration";
 export * from "./stores/PostgreSQLInviteLinkStore";
 export * from "./migrations/platform_settings_migration";
 export * from "./stores/PostgreSQLPlatformSettingsStore";
+export * from "./migrations/platform_audit_migration";
+export * from "./stores/PostgreSQLPlatformAuditStore";
 export * from "./migrations/migration";
 export * from "./migrations/thread_migrations";
 export * from "./migrations/thread_tenant_migration";

package/src/migrations/mcp_server_config_migrations.ts

@@ -53,3 +53,50 @@ export const createMcpServerConfigsTable: Migration = {
     await client.query('DROP TABLE IF EXISTS lattice_mcp_server_configs');
   },
 };
+
+/**
+ * Migration: create the personal MCP server config table.
+ *
+ * Personal configs are deliberately stored outside
+ * `lattice_mcp_server_configs` so tenant-only code paths (and older replicas
+ * during a rolling upgrade) can never read or register them.
+ */
+export const createPersonalMcpServerConfigsTable: Migration = {
+  version: 220,
+  name: 'create_personal_mcp_server_configs_table',
+  up: async (client: PoolClient) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_personal_mcp_server_configs (
+        id VARCHAR(255) NOT NULL,
+        tenant_id VARCHAR(255) NOT NULL,
+        owner_assistant_id VARCHAR(255) NOT NULL,
+        key VARCHAR(255) NOT NULL,
+        name VARCHAR(255),
+        description TEXT,
+        config JSONB NOT NULL,
+        selected_tools JSONB NOT NULL DEFAULT '[]',
+        is_env_encrypted BOOLEAN NOT NULL DEFAULT FALSE,
+        status VARCHAR(50) NOT NULL DEFAULT 'disconnected',
+        created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+        PRIMARY KEY (tenant_id, owner_assistant_id, id)
+      )
+    `);
+
+    await client.query(`
+      CREATE UNIQUE INDEX IF NOT EXISTS idx_lattice_personal_mcp_assistant_key
+      ON lattice_personal_mcp_server_configs(tenant_id, owner_assistant_id, key)
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_personal_mcp_assistant
+      ON lattice_personal_mcp_server_configs(tenant_id, owner_assistant_id)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query('DROP INDEX IF EXISTS idx_lattice_personal_mcp_assistant_key');
+    await client.query('DROP INDEX IF EXISTS idx_lattice_personal_mcp_assistant');
+    await client.query('DROP TABLE IF EXISTS lattice_personal_mcp_server_configs');
+  },
+};

package/src/migrations/platform_audit_migration.ts

@@ -0,0 +1,45 @@
+import type { Migration } from "./migration";
+
+/** Creates the append-only platform-level audit log table. */
+export const createPlatformAuditTable: Migration = {
+  version: 219,
+  name: "create_platform_audit_table",
+  up: async (client) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_platform_audit (
+        id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+        initiator_user_id VARCHAR(255),
+        initiator_role VARCHAR(16),
+        initiator JSONB NOT NULL,
+        source VARCHAR(32) NOT NULL,
+        executor JSONB NOT NULL,
+        tenant_id VARCHAR(255),
+        action VARCHAR(16) NOT NULL,
+        status VARCHAR(16) NOT NULL,
+        summary JSONB NOT NULL,
+        before JSONB,
+        after JSONB,
+        error_code VARCHAR(64),
+        duration_ms INTEGER NOT NULL,
+        created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+      )
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_platform_audit_created
+        ON lattice_platform_audit(created_at DESC, id DESC)
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_platform_audit_initiator
+        ON lattice_platform_audit(initiator_user_id, created_at DESC)
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_platform_audit_tenant
+        ON lattice_platform_audit(tenant_id, created_at DESC)
+        WHERE tenant_id IS NOT NULL
+    `);
+  },
+  down: async (client) => {
+    await client.query("DROP TABLE IF EXISTS lattice_platform_audit");
+  },
+};

package/src/stores/PostgreSQLMcpServerConfigStore.ts

@@ -7,12 +7,17 @@ import type { PoolConfig } from 'pg';
 import {
   McpServerConfigStore,
   McpServerConfigEntry,
+  PersonalMcpServerConfigStore,
+  PersonalMcpServerConfigEntry,
   CreateMcpServerConfigRequest,
   UpdateMcpServerConfigRequest,
   McpServerConfig,
 } from '@axiom-lattice/protocols';
 import { MigrationManager } from '../migrations/migration';
-import { createMcpServerConfigsTable } from '../migrations/mcp_server_config_migrations';
+import {
+  createMcpServerConfigsTable,
+  createPersonalMcpServerConfigsTable,
+} from '../migrations/mcp_server_config_migrations';
 import { encrypt, safeDecrypt } from '@axiom-lattice/core';
 
 /**
@@ -38,6 +43,26 @@ export interface PostgreSQLMcpServerConfigStoreOptions {
   autoMigrate?: boolean;
 }
 
+/** Row shape shared by the tenant and personal MCP config tables. */
+interface McpConfigRow {
+  id: string;
+  tenant_id: string;
+  key: string;
+  name: string | null;
+  description: string | null;
+  config: any;
+  selected_tools: string[];
+  is_env_encrypted: boolean;
+  status: string;
+  created_at: Date;
+  updated_at: Date;
+}
+
+/** Row shape for `lattice_personal_mcp_server_configs`. */
+interface PersonalMcpConfigRow extends McpConfigRow {
+  owner_assistant_id: string;
+}
+
 /**
  * PostgreSQL implementation of McpServerConfigStore
  *
@@ -46,7 +71,8 @@ export interface PostgreSQLMcpServerConfigStoreOptions {
  * - Automatic env encryption/decryption
  * - Unique constraint on (tenant_id, key)
  */
-export class PostgreSQLMcpServerConfigStore implements McpServerConfigStore {
+export class PostgreSQLMcpServerConfigStore
+  implements McpServerConfigStore, PersonalMcpServerConfigStore {
   private pool: Pool;
   private migrationManager!: MigrationManager;
   private initialized: boolean = false;
@@ -73,6 +99,7 @@ export class PostgreSQLMcpServerConfigStore implements McpServerConfigStore {
 
     this.migrationManager = new MigrationManager(this.pool);
     this.migrationManager.register(createMcpServerConfigsTable);
+    this.migrationManager.register(createPersonalMcpServerConfigsTable);
 
     // Auto-migrate by default
     if (options.autoMigrate !== false) {
@@ -411,6 +438,309 @@ export class PostgreSQLMcpServerConfigStore implements McpServerConfigStore {
     return result.rows.length > 0;
   }
 
+  /**
+   * Get all personal MCP server configurations owned by one assistant.
+   */
+  async getAllConfigsForAssistant(
+    tenantId: string,
+    ownerAssistantId: string,
+  ): Promise<PersonalMcpServerConfigEntry[]> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<PersonalMcpConfigRow>(
+      `
+      SELECT id, tenant_id, owner_assistant_id, key, name, description, config, selected_tools, is_env_encrypted, status, created_at, updated_at
+      FROM lattice_personal_mcp_server_configs
+      WHERE tenant_id = $1 AND owner_assistant_id = $2
+      ORDER BY created_at DESC
+    `,
+      [tenantId, ownerAssistantId],
+    );
+
+    return result.rows.map((row) => this.mapPersonalRowToEntry(row));
+  }
+
+  /**
+   * Get every personal MCP server configuration across tenants.
+   *
+   * Operations/tests only; startup restoration must not call this.
+   */
+  async getAllPersonalConfigsWithoutTenant(): Promise<PersonalMcpServerConfigEntry[]> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<PersonalMcpConfigRow>(
+      `
+      SELECT id, tenant_id, owner_assistant_id, key, name, description, config, selected_tools, is_env_encrypted, status, created_at, updated_at
+      FROM lattice_personal_mcp_server_configs
+      ORDER BY created_at DESC
+    `,
+    );
+
+    return result.rows.map((row) => this.mapPersonalRowToEntry(row));
+  }
+
+  /**
+   * Get one personal MCP server configuration by id.
+   */
+  async getConfigByIdForAssistant(
+    tenantId: string,
+    ownerAssistantId: string,
+    id: string,
+  ): Promise<PersonalMcpServerConfigEntry | null> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<PersonalMcpConfigRow>(
+      `
+      SELECT id, tenant_id, owner_assistant_id, key, name, description, config, selected_tools, is_env_encrypted, status, created_at, updated_at
+      FROM lattice_personal_mcp_server_configs
+      WHERE tenant_id = $1 AND owner_assistant_id = $2 AND id = $3
+    `,
+      [tenantId, ownerAssistantId, id],
+    );
+
+    if (result.rows.length === 0) {
+      return null;
+    }
+
+    return this.mapPersonalRowToEntry(result.rows[0]);
+  }
+
+  /**
+   * Get one personal MCP server configuration by business key.
+   */
+  async getConfigByKeyForAssistant(
+    tenantId: string,
+    ownerAssistantId: string,
+    key: string,
+  ): Promise<PersonalMcpServerConfigEntry | null> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<PersonalMcpConfigRow>(
+      `
+      SELECT id, tenant_id, owner_assistant_id, key, name, description, config, selected_tools, is_env_encrypted, status, created_at, updated_at
+      FROM lattice_personal_mcp_server_configs
+      WHERE tenant_id = $1 AND owner_assistant_id = $2 AND key = $3
+    `,
+      [tenantId, ownerAssistantId, key],
+    );
+
+    if (result.rows.length === 0) {
+      return null;
+    }
+
+    return this.mapPersonalRowToEntry(result.rows[0]);
+  }
+
+  /**
+   * Create a personal MCP server configuration.
+   */
+  async createConfigForAssistant(
+    tenantId: string,
+    ownerAssistantId: string,
+    id: string,
+    data: CreateMcpServerConfigRequest,
+  ): Promise<PersonalMcpServerConfigEntry> {
+    await this.ensureInitialized();
+
+    const now = new Date();
+    const nowString = now.toISOString();
+    const { config: configWithEncryptedEnv, isEnvEncrypted } =
+      this.encryptEnvInConfig(data.config);
+
+    await this.pool.query(
+      `
+      INSERT INTO lattice_personal_mcp_server_configs (id, tenant_id, owner_assistant_id, key, name, description, config, selected_tools, is_env_encrypted, status, created_at, updated_at)
+      VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11::timestamp, $12::timestamp)
+      ON CONFLICT (tenant_id, owner_assistant_id, id) DO UPDATE SET
+        key = EXCLUDED.key,
+        name = EXCLUDED.name,
+        description = EXCLUDED.description,
+        config = EXCLUDED.config,
+        selected_tools = EXCLUDED.selected_tools,
+        is_env_encrypted = EXCLUDED.is_env_encrypted,
+        status = EXCLUDED.status,
+        updated_at = EXCLUDED.updated_at
+    `,
+      [
+        id,
+        tenantId,
+        ownerAssistantId,
+        data.key,
+        data.name || null,
+        data.description || null,
+        JSON.stringify(configWithEncryptedEnv),
+        JSON.stringify(data.selectedTools || []),
+        isEnvEncrypted,
+        'disconnected',
+        nowString,
+        nowString,
+      ],
+    );
+
+    return {
+      id,
+      tenantId,
+      ownerAssistantId,
+      key: data.key,
+      config: data.config,
+      selectedTools: data.selectedTools || [],
+      isEnvEncrypted,
+      status: 'disconnected',
+      name: data.name,
+      description: data.description,
+      createdAt: now,
+      updatedAt: now,
+    };
+  }
+
+  /**
+   * Update a personal MCP server configuration.
+   */
+  async updateConfigForAssistant(
+    tenantId: string,
+    ownerAssistantId: string,
+    id: string,
+    updates: Partial<UpdateMcpServerConfigRequest>,
+  ): Promise<PersonalMcpServerConfigEntry | null> {
+    await this.ensureInitialized();
+
+    const existing = await this.getConfigByIdForAssistant(
+      tenantId,
+      ownerAssistantId,
+      id,
+    );
+    if (!existing) {
+      return null;
+    }
+
+    const updateData: Record<string, any> = {};
+
+    if (updates.key !== undefined) {
+      updateData.key = updates.key;
+    }
+
+    if (updates.name !== undefined) {
+      updateData.name = updates.name || null;
+    }
+
+    if (updates.description !== undefined) {
+      updateData.description = updates.description || null;
+    }
+
+    if (updates.config !== undefined) {
+      const { config: configWithEncryptedEnv, isEnvEncrypted } =
+        this.encryptEnvInConfig(updates.config);
+      updateData.config = JSON.stringify(configWithEncryptedEnv);
+      updateData.is_env_encrypted = isEnvEncrypted;
+    }
+
+    if (updates.selectedTools !== undefined) {
+      updateData.selected_tools = JSON.stringify(updates.selectedTools);
+    }
+
+    if (updates.status !== undefined) {
+      updateData.status = updates.status;
+    }
+
+    if (Object.keys(updateData).length === 0) {
+      return existing;
+    }
+
+    updateData.updated_at = new Date().toISOString();
+
+    const fields = Object.keys(updateData);
+    const values = Object.values(updateData);
+    values.push(tenantId);
+    values.push(ownerAssistantId);
+    values.push(id);
+
+    const setClauses = fields.map((field, index) => {
+      if (field === 'updated_at') {
+        return `${field} = $${index + 1}::timestamp`;
+      }
+      return `${field} = $${index + 1}`;
+    });
+
+    const whereTenantIndex = fields.length + 1;
+    const whereOwnerIndex = fields.length + 2;
+    const whereIdIndex = fields.length + 3;
+
+    const sql = `
+      UPDATE lattice_personal_mcp_server_configs
+      SET ${setClauses.join(', ')}
+      WHERE tenant_id = $${whereTenantIndex}
+        AND owner_assistant_id = $${whereOwnerIndex}
+        AND id = $${whereIdIndex}
+    `;
+
+    await this.pool.query(sql, values);
+
+    return this.getConfigByIdForAssistant(tenantId, ownerAssistantId, id);
+  }
+
+  /**
+   * Delete one personal MCP server configuration by id.
+   */
+  async deleteConfigForAssistant(
+    tenantId: string,
+    ownerAssistantId: string,
+    id: string,
+  ): Promise<boolean> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query(
+      `
+      DELETE FROM lattice_personal_mcp_server_configs
+      WHERE tenant_id = $1 AND owner_assistant_id = $2 AND id = $3
+    `,
+      [tenantId, ownerAssistantId, id],
+    );
+
+    return result.rowCount !== null && result.rowCount > 0;
+  }
+
+  /**
+   * Delete every personal MCP server configuration owned by one assistant.
+   */
+  async deleteConfigsForAssistant(
+    tenantId: string,
+    ownerAssistantId: string,
+  ): Promise<number> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query(
+      `
+      DELETE FROM lattice_personal_mcp_server_configs
+      WHERE tenant_id = $1 AND owner_assistant_id = $2
+    `,
+      [tenantId, ownerAssistantId],
+    );
+
+    return result.rowCount ?? 0;
+  }
+
+  /**
+   * Check whether a personal MCP server configuration exists.
+   */
+  async hasConfigForAssistant(
+    tenantId: string,
+    ownerAssistantId: string,
+    id: string,
+  ): Promise<boolean> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query(
+      `
+      SELECT 1 FROM lattice_personal_mcp_server_configs
+      WHERE tenant_id = $1 AND owner_assistant_id = $2 AND id = $3
+      LIMIT 1
+    `,
+      [tenantId, ownerAssistantId, id],
+    );
+
+    return result.rows.length > 0;
+  }
+
   /**
    * Dispose resources and close the connection pool
    */
@@ -433,19 +763,7 @@ export class PostgreSQLMcpServerConfigStore implements McpServerConfigStore {
    * Map database row to McpServerConfigEntry object
    * Automatically decrypts env values if present
    */
-  private mapRowToEntry(row: {
-    id: string;
-    tenant_id: string;
-    key: string;
-    name: string | null;
-    description: string | null;
-    config: any;
-    selected_tools: string[];
-    is_env_encrypted: boolean;
-    status: string;
-    created_at: Date;
-    updated_at: Date;
-  }): McpServerConfigEntry {
+  private mapRowToEntry(row: McpConfigRow): McpServerConfigEntry {
     const config: McpServerConfig = typeof row.config === 'string'
       ? JSON.parse(row.config)
       : row.config;
@@ -483,6 +801,16 @@ export class PostgreSQLMcpServerConfigStore implements McpServerConfigStore {
     };
   }
 
+  /**
+   * Map a personal table row, tagging the owning assistant.
+   */
+  private mapPersonalRowToEntry(row: PersonalMcpConfigRow): PersonalMcpServerConfigEntry {
+    return {
+      ...this.mapRowToEntry(row),
+      ownerAssistantId: row.owner_assistant_id,
+    };
+  }
+
   /**
    * Encrypt env values in config before storing
    */