@axiom-lattice/pg-stores 3.1.6 → 3.1.9

package/src/migrations/create_open_audit_table.ts

@@ -0,0 +1,36 @@
+import type { Migration } from "./migration";
+import type { PoolClient } from "pg";
+
+/**
+ * Append-only audit log for the Open API surface (one row per invoke).
+ */
+export const createOpenAuditTable: Migration = {
+  version: 179,
+  name: "create_open_audit_table",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_open_audit (
+        id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+        tenant_id VARCHAR(255) NOT NULL,
+        credential_id VARCHAR(255) NOT NULL,
+        credential_kind VARCHAR(32) NOT NULL,
+        project_id VARCHAR(255) NOT NULL,
+        domain VARCHAR(255) NOT NULL,
+        action VARCHAR(255) NOT NULL,
+        args JSONB,
+        status VARCHAR(16) NOT NULL,
+        error_code VARCHAR(64),
+        duration_ms INTEGER NOT NULL,
+        created_at TIMESTAMP NOT NULL DEFAULT NOW()
+      )
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_open_audit_tenant_created
+      ON lattice_open_audit(tenant_id, created_at DESC)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query("DROP TABLE IF EXISTS lattice_open_audit");
+  },
+};

package/src/migrations/project_room_migration.ts

@@ -126,3 +126,29 @@ export const createProjectRoomTables: Migration = {
     await client.query("DROP TABLE IF EXISTS lattice_project_rooms");
   },
 };
+
+/** Adds per-user, per-room read markers for unread counts. */
+export const createProjectRoomReadStates: Migration = {
+  version: 177,
+  name: "create_project_room_read_states",
+  up: async (client) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_project_room_read_states (
+        tenant_id VARCHAR(255) NOT NULL,
+        room_id VARCHAR(255) NOT NULL,
+        user_id VARCHAR(255) NOT NULL,
+        last_read_at TIMESTAMPTZ NOT NULL,
+        updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+        PRIMARY KEY (tenant_id, room_id, user_id)
+      )
+    `);
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_project_room_read_states_user
+      ON lattice_project_room_read_states (tenant_id, user_id)
+    `);
+  },
+  down: async (client) => {
+    await client.query("DROP INDEX IF EXISTS idx_lattice_project_room_read_states_user");
+    await client.query("DROP TABLE IF EXISTS lattice_project_room_read_states");
+  },
+};

package/src/stores/PostgreSQLA2AApiKeyStore.ts

@@ -1,6 +1,7 @@
 import { Pool } from "pg";
 import type { PoolConfig } from "pg";
 import type {
+  UpdateA2AApiKeyInput,
   A2AApiKeyStore,
   A2AApiKeyRecord,
   CreateA2AApiKeyInput,
@@ -9,6 +10,7 @@ import type {
 import { MigrationManager } from "../migrations/migration";
 import { createA2AApiKeysTable } from "../migrations/a2a_api_key_migration";
 import { addA2AKeyAssistantIds } from "../migrations/add_a2a_key_assistant_ids";
+import { addApiKeyGrants } from "../migrations/add_api_key_grants";
 import { encrypt, decrypt } from "@axiom-lattice/core";
 import { randomUUID } from "crypto";
 
@@ -24,6 +26,7 @@ type KeyRow = {
   tenant_id: string;
   project_id: string;
   assistant_ids: string[] | null;
+  grants: unknown[] | null;
   label: string | null;
   enabled: boolean;
   created_at: Date;
@@ -41,6 +44,7 @@ function mapRowToRecord(row: KeyRow): A2AApiKeyRecord {
     tenantId: row.tenant_id,
     projectId: row.project_id,
     assistantIds: row.assistant_ids ?? undefined,
+    grants: (row.grants as A2AApiKeyRecord["grants"]) ?? undefined,
     label: row.label || undefined,
     enabled: row.enabled,
     createdAt: row.created_at,
@@ -78,6 +82,7 @@ export class PostgreSQLA2AApiKeyStore implements A2AApiKeyStore {
     // overlap is safe.
     this.migrationManager.register(createA2AApiKeysTable);
     this.migrationManager.register(addA2AKeyAssistantIds);
+    this.migrationManager.register(addApiKeyGrants);
 
     if (options.autoMigrate !== false) {
       this.initialize().catch((error) => {
@@ -111,13 +116,31 @@ export class PostgreSQLA2AApiKeyStore implements A2AApiKeyStore {
     if (!this.initialized) await this.initialize();
   }
 
+  /** Map rows that decrypt cleanly; skip rows under an unknown key. */
+  private safeMap(rows: KeyRow[]): A2AApiKeyRecord[] {
+    const out: A2AApiKeyRecord[] = [];
+    for (const row of rows) {
+      try {
+        out.push(mapRowToRecord(row));
+      } catch {
+        // Skip undecryptable rows.
+      }
+    }
+    return out;
+  }
+
   async findByKey(key: string): Promise<A2AApiKeyRecord | null> {
     await this.ensureInitialized();
     const all = await this.pool.query<KeyRow>(
       `SELECT * FROM lattice_a2a_api_keys WHERE enabled = true`,
     );
     for (const row of all.rows) {
-      if (decrypt(row.key_value) === key) return mapRowToRecord(row);
+      // A row encrypted under a different key must not break every lookup.
+      try {
+        if (decrypt(row.key_value) === key) return mapRowToRecord(row);
+      } catch {
+        continue;
+      }
     }
     return null;
   }
@@ -128,7 +151,12 @@ export class PostgreSQLA2AApiKeyStore implements A2AApiKeyStore {
       `SELECT * FROM lattice_a2a_api_keys WHERE id = $1`,
       [id],
     );
-    return result.rows[0] ? mapRowToRecord(result.rows[0]) : null;
+    if (!result.rows[0]) return null;
+    try {
+      return mapRowToRecord(result.rows[0]);
+    } catch {
+      return null;
+    }
   }
 
   async list(params: { tenantId?: string; limit?: number; offset?: number }): Promise<A2AApiKeyRecord[]> {
@@ -140,28 +168,71 @@ export class PostgreSQLA2AApiKeyStore implements A2AApiKeyStore {
         `SELECT * FROM lattice_a2a_api_keys WHERE tenant_id = $1 ORDER BY created_at DESC LIMIT $2 OFFSET $3`,
         [params.tenantId, limit, offset],
       );
-      return result.rows.map(mapRowToRecord);
+      return this.safeMap(result.rows);
     }
     const result = await this.pool.query<KeyRow>(
       `SELECT * FROM lattice_a2a_api_keys ORDER BY created_at DESC LIMIT $1 OFFSET $2`,
       [limit, offset],
     );
-    return result.rows.map(mapRowToRecord);
+    return this.safeMap(result.rows);
   }
 
   async create(input: CreateA2AApiKeyInput): Promise<A2AApiKeyRecord> {
     await this.ensureInitialized();
     const key = generateApiKey();
     const result = await this.pool.query<KeyRow>(
-      `INSERT INTO lattice_a2a_api_keys (key_value, tenant_id, project_id, assistant_ids, label)
-       VALUES ($1, $2, $3, $4, $5) RETURNING *`,
-      [encrypt(key), input.tenantId, input.projectId, input.assistantIds ? JSON.stringify(input.assistantIds) : null, input.label || null],
+      `INSERT INTO lattice_a2a_api_keys (key_value, tenant_id, project_id, assistant_ids, grants, label)
+       VALUES ($1, $2, $3, $4, $5, $6) RETURNING *`,
+      [
+        encrypt(key),
+        input.tenantId,
+        input.projectId,
+        input.assistantIds ? JSON.stringify(input.assistantIds) : null,
+        input.grants ? JSON.stringify(input.grants) : null,
+        input.label || null,
+      ],
     );
     const record = mapRowToRecord(result.rows[0]);
     record.key = key;
     return record;
   }
 
+
+  async update(id: string, input: UpdateA2AApiKeyInput): Promise<A2AApiKeyRecord> {
+    await this.ensureInitialized();
+    const sets: string[] = [];
+    const values: unknown[] = [];
+    if (input.label !== undefined) {
+      values.push(input.label || null);
+      sets.push(`label = $${values.length}`);
+    }
+    if (input.projectId !== undefined) {
+      values.push(input.projectId);
+      sets.push(`project_id = $${values.length}`);
+    }
+    if (input.assistantIds !== undefined) {
+      values.push(input.assistantIds ? JSON.stringify(input.assistantIds) : null);
+      sets.push(`assistant_ids = $${values.length}`);
+    }
+    if (input.grants !== undefined) {
+      values.push(input.grants ? JSON.stringify(input.grants) : null);
+      sets.push(`grants = $${values.length}`);
+    }
+    if (sets.length === 0) {
+      const existing = await this.findById(id);
+      if (!existing) throw new Error(`A2A API key not found: ${id}`);
+      return existing;
+    }
+    sets.push("updated_at = NOW()");
+    values.push(id);
+    const result = await this.pool.query<KeyRow>(
+      `UPDATE lattice_a2a_api_keys SET ${sets.join(", ")} WHERE id = $${values.length} RETURNING *`,
+      values,
+    );
+    if (result.rows.length === 0) throw new Error(`A2A API key not found: ${id}`);
+    return mapRowToRecord(result.rows[0]);
+  }
+
   async disable(id: string): Promise<A2AApiKeyRecord> {
     await this.ensureInitialized();
     const result = await this.pool.query<KeyRow>(
@@ -207,13 +278,18 @@ export class PostgreSQLA2AApiKeyStore implements A2AApiKeyStore {
     );
     const map = new Map<string, A2AApiKeyEntry>();
     for (const row of result.rows) {
-      const key = decrypt(row.key_value);
-      map.set(key, {
-        key,
-        tenantId: row.tenant_id,
-        projectId: row.project_id,
-        assistantIds: row.assistant_ids ?? undefined,
-      });
+      try {
+        const key = decrypt(row.key_value);
+        map.set(key, {
+          key,
+          tenantId: row.tenant_id,
+          projectId: row.project_id,
+          assistantIds: row.assistant_ids ?? undefined,
+          grants: (row.grants as A2AApiKeyEntry["grants"]) ?? undefined,
+        });
+      } catch {
+        // Skip rows encrypted under an unknown key.
+      }
     }
     return map;
   }

package/src/stores/PostgreSQLMcpServerConfigStore.ts

@@ -464,6 +464,20 @@ export class PostgreSQLMcpServerConfigStore implements McpServerConfigStore {
       }
     }
 
+    // Decrypt header values if present and encrypted
+    if (config.headers && row.is_env_encrypted) {
+      try {
+        const decryptedHeaders: Record<string, string> = {};
+        for (const [key, value] of Object.entries(config.headers)) {
+          decryptedHeaders[key] = decrypt(value);
+        }
+        config.headers = decryptedHeaders;
+      } catch (error) {
+        console.error('Failed to decrypt MCP server headers:', error);
+        throw new Error('Failed to decrypt MCP server configuration');
+      }
+    }
+
     return {
       id: row.id,
       tenantId: row.tenant_id,
@@ -495,6 +509,15 @@ export class PostgreSQLMcpServerConfigStore implements McpServerConfigStore {
       isEnvEncrypted = true;
     }
 
+    if (configCopy.headers && Object.keys(configCopy.headers).length > 0) {
+      const encryptedHeaders: Record<string, string> = {};
+      for (const [key, value] of Object.entries(configCopy.headers)) {
+        encryptedHeaders[key] = encrypt(value);
+      }
+      configCopy.headers = encryptedHeaders;
+      isEnvEncrypted = true;
+    }
+
     return { config: configCopy, isEnvEncrypted };
   }
 }

package/src/stores/PostgreSQLOpenAuditStore.ts

@@ -0,0 +1,143 @@
+import { Pool } from "pg";
+import type { PoolConfig } from "pg";
+import type {
+  OpenAuditStore,
+  OpenAuditAppendInput,
+  OpenAuditQuery,
+  OpenAuditRecord,
+} from "@axiom-lattice/protocols";
+import { MigrationManager } from "../migrations/migration";
+import { createOpenAuditTable } from "../migrations/create_open_audit_table";
+
+export interface PostgreSQLOpenAuditStoreOptions {
+  pool?: Pool;
+  poolConfig?: string | PoolConfig;
+  autoMigrate?: boolean;
+}
+
+/** Append-only PostgreSQL audit store for the Open API surface. */
+export class PostgreSQLOpenAuditStore implements OpenAuditStore {
+  private pool: Pool;
+  private migrationManager!: MigrationManager;
+  private initialized = false;
+  private ownsPool = true;
+  private initPromise: Promise<void> | null = null;
+
+  constructor(options: PostgreSQLOpenAuditStoreOptions) {
+    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 as PoolConfig)
+          : (() => {
+              throw new Error("Either pool or poolConfig must be provided");
+            })();
+    this.migrationManager = new MigrationManager(this.pool);
+    this.migrationManager.register(createOpenAuditTable);
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error("Failed to initialize PostgreSQLOpenAuditStore:", error);
+        throw error;
+      });
+    }
+  }
+
+  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;
+  }
+
+  async dispose(): Promise<void> {
+    if (this.ownsPool && this.pool) await this.pool.end();
+  }
+
+  private async ensureInitialized(): Promise<void> {
+    if (!this.initialized) await this.initialize();
+  }
+
+  async query(params: OpenAuditQuery): Promise<OpenAuditRecord[]> {
+    await this.ensureInitialized();
+    const clauses = ["tenant_id = $1"];
+    const values: unknown[] = [params.tenantId];
+    if (params.credentialId) {
+      values.push(params.credentialId);
+      clauses.push(`credential_id = $${values.length}`);
+    }
+    if (params.domain) {
+      values.push(params.domain);
+      clauses.push(`domain = $${values.length}`);
+    }
+    if (params.status) {
+      values.push(params.status);
+      clauses.push(`status = $${values.length}`);
+    }
+    if (params.from) {
+      values.push(params.from);
+      clauses.push(`created_at >= $${values.length}`);
+    }
+    if (params.to) {
+      values.push(params.to);
+      clauses.push(`created_at < $${values.length}`);
+    }
+    const limit = params.limit ?? 100;
+    const offset = params.offset ?? 0;
+    values.push(limit, offset);
+    const { rows } = await this.pool.query(
+      `SELECT * FROM lattice_open_audit
+       WHERE ${clauses.join(" AND ")}
+       ORDER BY created_at DESC
+       LIMIT $${values.length - 1} OFFSET $${values.length}`,
+      values,
+    );
+    return rows.map((row) => ({
+      id: row.id,
+      tenantId: row.tenant_id,
+      credentialId: row.credential_id,
+      credentialKind: row.credential_kind,
+      projectId: row.project_id,
+      domain: row.domain,
+      action: row.action,
+      args: row.args,
+      status: row.status,
+      errorCode: row.error_code ?? undefined,
+      durationMs: row.duration_ms,
+      createdAt: row.created_at,
+    }));
+  }
+
+  async append(record: OpenAuditAppendInput): Promise<void> {
+    await this.ensureInitialized();
+    await this.pool.query(
+      `INSERT INTO lattice_open_audit
+        (tenant_id, credential_id, credential_kind, project_id, domain, action, args, status, error_code, duration_ms)
+       VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10)`,
+      [
+        record.tenantId,
+        record.credentialId,
+        record.credentialKind,
+        record.projectId,
+        record.domain,
+        record.action,
+        record.args === undefined ? null : JSON.stringify(record.args),
+        record.status,
+        record.errorCode ?? null,
+        Math.max(0, Math.round(record.durationMs)),
+      ],
+    );
+  }
+}

package/src/stores/PostgreSQLProjectRoomMessageStore.ts

@@ -83,6 +83,25 @@ function mapRow(row: MessageRow): ProjectRoomMessage {
     ...(row.idempotency_key === null ? {} : { idempotencyKey: row.idempotency_key }), createdAt: row.created_at,
   };
 }
+function mapRowsTolerant(
+  rows: MessageRow[],
+  scope: { tenantId: string; roomId: string },
+): ProjectRoomMessage[] {
+  const messages: ProjectRoomMessage[] = [];
+  for (const row of rows) {
+    try {
+      messages.push(mapRow(row));
+    } catch (error) {
+      console.warn("[ProjectRoomMessageStore] Skipping invalid message row", {
+        tenantId: scope.tenantId,
+        roomId: scope.roomId,
+        rowId: typeof row.id === "string" ? row.id : undefined,
+        reason: error instanceof Error ? error.message : String(error),
+      });
+    }
+  }
+  return messages;
+}
 function isUniqueViolation(error: unknown): error is PostgreSQLError {
   return typeof error === "object" && error !== null && "code" in error
     && (error as PostgreSQLError).code === "23505";
@@ -159,7 +178,7 @@ export class PostgreSQLProjectRoomMessageStore implements ProjectRoomMessageStor
          WHERE tenant_id = $1 AND room_id = $2 ORDER BY created_at DESC, id DESC LIMIT $3`,
         [input.tenantId, input.roomId, limit],
       );
-    return result.rows.map(mapRow);
+    return mapRowsTolerant(result.rows, { tenantId: input.tenantId, roomId: input.roomId });
   }
 
   /** Finds a message by tenant-scoped ID. */
@@ -170,6 +189,22 @@ export class PostgreSQLProjectRoomMessageStore implements ProjectRoomMessageStor
     return result.rows[0] ? mapRow(result.rows[0]) : null;
   }
 
+  /** Counts messages created strictly after a horizon, optionally excluding one human author. */
+  async countAfter(input: {
+    tenantId: string; roomId: string; after: Date; excludeAuthorUserId?: string;
+  }): Promise<number> {
+    if (!(input.after instanceof Date) || Number.isNaN(input.after.getTime())) {
+      throw new RangeError("Project room message count horizon date is invalid");
+    }
+    const result = await this.pool.query<{ count: number }>(
+      `SELECT COUNT(*)::int AS count FROM lattice_project_room_messages
+       WHERE tenant_id = $1 AND room_id = $2 AND created_at > $3
+         AND ($4::text IS NULL OR author->>'type' <> 'human' OR author->>'userId' <> $4)`,
+      [input.tenantId, input.roomId, input.after, input.excludeAuthorUserId ?? null],
+    );
+    return result.rows[0]?.count ?? 0;
+  }
+
   private parameters(input: Omit<ProjectRoomMessage, "createdAt">): unknown[] {
     return [input.id, input.tenantId, input.workspaceId, input.projectId, input.roomId, JSON.stringify(input.author), JSON.stringify(input.content),
       JSON.stringify(input.mentions), input.replyToMessageId ?? null, input.source, input.sourceId ?? null, input.idempotencyKey ?? null];

package/src/stores/PostgreSQLProjectRoomReadStateStore.ts

@@ -0,0 +1,47 @@
+import type { ProjectRoomReadState, ProjectRoomReadStateStore } from "@axiom-lattice/protocols";
+import type { Pool } from "pg";
+
+type ReadStateRow = { tenant_id: unknown; room_id: unknown; user_id: unknown; last_read_at: unknown; updated_at: unknown };
+
+function mapRow(row: ReadStateRow): ProjectRoomReadState {
+  if (typeof row.tenant_id !== "string" || typeof row.room_id !== "string" || typeof row.user_id !== "string"
+    || !(row.last_read_at instanceof Date) || Number.isNaN(row.last_read_at.getTime())
+    || !(row.updated_at instanceof Date) || Number.isNaN(row.updated_at.getTime())) {
+    throw new Error("Invalid project room read state row");
+  }
+  return { tenantId: row.tenant_id, roomId: row.room_id, userId: row.user_id, lastReadAt: row.last_read_at, updatedAt: row.updated_at };
+}
+
+/** Persists tenant-isolated per-user project room read markers in PostgreSQL. */
+export class PostgreSQLProjectRoomReadStateStore implements ProjectRoomReadStateStore {
+  /** Creates a store using an externally managed shared pool. */
+  constructor(options: { pool: Pool }) { this.pool = options.pool; }
+  private readonly pool: Pool;
+
+  /** Returns the user's read marker for a room, or null when never reported. */
+  async get(tenantId: string, roomId: string, userId: string): Promise<ProjectRoomReadState | null> {
+    const result = await this.pool.query<ReadStateRow>(
+      `SELECT tenant_id, room_id, user_id, last_read_at, updated_at
+       FROM lattice_project_room_read_states WHERE tenant_id = $1 AND room_id = $2 AND user_id = $3`,
+      [tenantId, roomId, userId],
+    );
+    return result.rows[0] ? mapRow(result.rows[0]) : null;
+  }
+
+  /** Upserts the marker monotonically using GREATEST so older reports never regress it. */
+  async markRead(input: { tenantId: string; roomId: string; userId: string; lastReadAt: Date }): Promise<ProjectRoomReadState> {
+    if (!(input.lastReadAt instanceof Date) || Number.isNaN(input.lastReadAt.getTime())) {
+      throw new RangeError("Project room read state lastReadAt is invalid");
+    }
+    const result = await this.pool.query<ReadStateRow>(
+      `INSERT INTO lattice_project_room_read_states (tenant_id, room_id, user_id, last_read_at, updated_at)
+       VALUES ($1, $2, $3, $4, NOW())
+       ON CONFLICT (tenant_id, room_id, user_id) DO UPDATE
+         SET last_read_at = GREATEST(lattice_project_room_read_states.last_read_at, EXCLUDED.last_read_at),
+             updated_at = NOW()
+       RETURNING tenant_id, room_id, user_id, last_read_at, updated_at`,
+      [input.tenantId, input.roomId, input.userId, input.lastReadAt],
+    );
+    return mapRow(result.rows[0]);
+  }
+}

package/src/stores/PostgresSharedResourceStore.ts

@@ -25,6 +25,7 @@ type ShareRow = {
   password_hash: string | null;
   expires_at: string | null;
   max_access: number | null;
+  open_grants: unknown[] | null;
   access_count: number;
   revoked: boolean;
   created_at: string;
@@ -106,8 +107,8 @@ export class PostgresSharedResourceStore implements SharedResourceStore {
       `INSERT INTO lattice_shared_resources
        (token, address, title, created_by, tenant_id, workspace_id, project_id,
         assistant_id, vm_isolation, visibility, password_hash,
-        expires_at, max_access, revoked)
-       VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14)
+        expires_at, max_access, open_grants, revoked)
+       VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15)
        RETURNING *`,
       [
         record.token,
@@ -123,6 +124,7 @@ export class PostgresSharedResourceStore implements SharedResourceStore {
         record.passwordHash,
         record.expiresAt ? record.expiresAt.toISOString() : null,
         record.maxAccess,
+        record.openGrants ? JSON.stringify(record.openGrants) : null,
         record.revoked,
       ],
     );
@@ -229,6 +231,7 @@ export class PostgresSharedResourceStore implements SharedResourceStore {
       passwordHash: row.password_hash,
       expiresAt: row.expires_at ? new Date(row.expires_at) : null,
       maxAccess: row.max_access,
+      openGrants: (row.open_grants as ShareRecord["openGrants"]) ?? undefined,
       accessCount: row.access_count ?? 0,
       revoked: row.revoked ?? false,
       createdAt: new Date(row.created_at),