@axiom-lattice/pg-stores 3.1.6 → 3.1.8

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "3.1.6",
+  "version": "3.1.8",
   "description": "PG stores implementation for Axiom Lattice framework",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -25,8 +25,8 @@
     "@langchain/core": "1.1.30",
     "pg": "^8.16.3",
     "uuid": "^9.0.1",
-    "@axiom-lattice/core": "4.2.6",
-    "@axiom-lattice/protocols": "4.1.4"
+    "@axiom-lattice/core": "4.2.8",
+    "@axiom-lattice/protocols": "4.2.0"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",

package/src/__tests__/PostgreSQLMcpServerConfigStore.test.ts

@@ -0,0 +1,79 @@
+import { describe, expect, it } from "@jest/globals";
+import type { Pool } from "pg";
+
+// The @axiom-lattice/core barrel cannot load under jest in this package (ESM-only
+// transitive deps), so mock it with the real encryption module — node:crypto only.
+jest.mock("@axiom-lattice/core", () => ({
+  ...jest.requireActual("../../../core/src/util/encryption"),
+}));
+
+import { PostgreSQLMcpServerConfigStore } from "../stores/PostgreSQLMcpServerConfigStore";
+
+const makeRow = (configJson: string, isEnvEncrypted: boolean) => ({
+  id: "cfg-1",
+  tenant_id: "tenant-1",
+  key: "remote",
+  name: null,
+  description: null,
+  config: configJson,
+  selected_tools: [],
+  is_env_encrypted: isEnvEncrypted,
+  status: "disconnected",
+  created_at: new Date(),
+  updated_at: new Date(),
+});
+
+describe("PostgreSQLMcpServerConfigStore header encryption", () => {
+  it("encrypts header values at rest and decrypts on read", async () => {
+    let storedConfig = "";
+    let storedFlag = false;
+    const query = jest.fn(async (sql: string, values: unknown[]) => {
+      if (sql.includes("INSERT INTO lattice_mcp_server_configs")) {
+        storedConfig = values[5] as string;
+        storedFlag = values[7] as boolean;
+        return { rowCount: 1 };
+      }
+      if (sql.includes("SELECT id, tenant_id, key")) {
+        return { rows: [makeRow(storedConfig, storedFlag)] };
+      }
+      return { rows: [] };
+    });
+    const store = new PostgreSQLMcpServerConfigStore({ pool: { query } as unknown as Pool });
+
+    await store.createConfig("tenant-1", "cfg-1", {
+      key: "remote",
+      config: {
+        name: "remote",
+        version: "1.0.0",
+        transport: "streamable_http",
+        url: "https://mcp.example.com/mcp",
+        headers: { Authorization: "Bearer sk-secret-token" },
+      },
+    });
+
+    expect(storedFlag).toBe(true);
+    const atRest = JSON.parse(storedConfig);
+    expect(atRest.headers.Authorization).not.toBe("Bearer sk-secret-token");
+
+    const read = await store.getConfigByKey("tenant-1", "remote");
+    expect(read?.config.headers?.Authorization).toBe("Bearer sk-secret-token");
+    expect(read?.config.headers?.Authorization).not.toBe(atRest.headers.Authorization);
+  });
+
+  it("leaves legacy rows without headers untouched", async () => {
+    const legacyConfig = JSON.stringify({
+      name: "legacy", version: "1.0.0", transport: "stdio", command: "npx",
+    });
+    const query = jest.fn(async (sql: string) => {
+      if (sql.includes("SELECT id, tenant_id, key")) {
+        return { rows: [makeRow(legacyConfig, false)] };
+      }
+      return { rows: [] };
+    });
+    const store = new PostgreSQLMcpServerConfigStore({ pool: { query } as unknown as Pool });
+
+    const read = await store.getConfigByKey("tenant-1", "remote");
+    expect(read?.config.headers).toBeUndefined();
+    expect(read?.config.command).toBe("npx");
+  });
+});

package/src/__tests__/PostgreSQLProjectRoomReadStateStore.test.ts

@@ -0,0 +1,136 @@
+import { randomUUID } from "node:crypto";
+import { afterAll, beforeAll, beforeEach, describe, expect, it, jest } from "@jest/globals";
+import { Pool } from "pg";
+import type { ProjectRoomMessage } from "@axiom-lattice/protocols";
+
+import { MigrationManager } from "../migrations/migration";
+import { createProjectRoomReadStates, createProjectRoomTables } from "../migrations/project_room_migration";
+import { PostgreSQLProjectRoomMessageStore } from "../stores/PostgreSQLProjectRoomMessageStore";
+import { PostgreSQLProjectRoomReadStateStore } from "../stores/PostgreSQLProjectRoomReadStateStore";
+
+const query = jest.fn();
+const mockPool = { query } as unknown as Pool;
+const readStateRow = {
+  tenant_id: "tenant-1", room_id: "room-1", user_id: "user-1",
+  last_read_at: new Date("2026-01-02T00:00:00.000Z"), updated_at: new Date("2026-01-02T00:00:01.000Z"),
+};
+
+describe("PostgreSQLProjectRoomReadStateStore SQL contract", () => {
+  beforeEach(() => query.mockReset());
+
+  it("returns null when no marker exists and selects by the full primary key", async () => {
+    query.mockResolvedValueOnce({ rows: [] });
+    const store = new PostgreSQLProjectRoomReadStateStore({ pool: mockPool });
+    await expect(store.get("tenant-1", "room-1", "user-1")).resolves.toBeNull();
+    expect(query.mock.calls[0][0]).toContain("FROM lattice_project_room_read_states");
+    expect(query.mock.calls[0][0]).toContain("WHERE tenant_id = $1 AND room_id = $2 AND user_id = $3");
+    expect(query.mock.calls[0][1]).toEqual(["tenant-1", "room-1", "user-1"]);
+  });
+
+  it("upserts monotonically with GREATEST and returns the stored row", async () => {
+    query.mockResolvedValueOnce({ rows: [readStateRow] });
+    const store = new PostgreSQLProjectRoomReadStateStore({ pool: mockPool });
+    const lastReadAt = new Date("2026-01-02T00:00:00.000Z");
+    await expect(store.markRead({ tenantId: "tenant-1", roomId: "room-1", userId: "user-1", lastReadAt }))
+      .resolves.toEqual({
+        tenantId: "tenant-1", roomId: "room-1", userId: "user-1",
+        lastReadAt, updatedAt: new Date("2026-01-02T00:00:01.000Z"),
+      });
+    const [sql, params] = query.mock.calls[0] as [string, unknown[]];
+    expect(sql).toContain("ON CONFLICT (tenant_id, room_id, user_id) DO UPDATE");
+    expect(sql).toContain("GREATEST(lattice_project_room_read_states.last_read_at, EXCLUDED.last_read_at)");
+    expect(params).toEqual(["tenant-1", "room-1", "user-1", lastReadAt]);
+  });
+
+  it("rejects an invalid markRead horizon before querying", async () => {
+    const store = new PostgreSQLProjectRoomReadStateStore({ pool: mockPool });
+    await expect(store.markRead({ tenantId: "tenant-1", roomId: "room-1", userId: "user-1", lastReadAt: new Date("bad") }))
+      .rejects.toThrow(RangeError);
+    expect(query).not.toHaveBeenCalled();
+  });
+
+  it("rejects malformed read state rows", async () => {
+    query.mockResolvedValueOnce({ rows: [{ ...readStateRow, last_read_at: "2026-01-02" }] });
+    const store = new PostgreSQLProjectRoomReadStateStore({ pool: mockPool });
+    await expect(store.get("tenant-1", "room-1", "user-1")).rejects.toThrow("Invalid project room read state row");
+  });
+});
+
+describe("PostgreSQLProjectRoomMessageStore countAfter SQL contract", () => {
+  beforeEach(() => query.mockReset());
+
+  it("counts strictly after the horizon and excludes one human author", async () => {
+    query.mockResolvedValueOnce({ rows: [{ count: 2 }] });
+    const store = new PostgreSQLProjectRoomMessageStore({ pool: mockPool });
+    const after = new Date("2026-01-01T00:00:00.000Z");
+    await expect(store.countAfter({ tenantId: "tenant-1", roomId: "room-1", after, excludeAuthorUserId: "user-1" }))
+      .resolves.toBe(2);
+    const [sql, params] = query.mock.calls[0] as [string, unknown[]];
+    expect(sql).toContain("created_at > $3");
+    expect(sql).toContain("($4::text IS NULL OR author->>'type' <> 'human' OR author->>'userId' <> $4)");
+    expect(params).toEqual(["tenant-1", "room-1", after, "user-1"]);
+  });
+
+  it("passes null when no author is excluded and rejects invalid horizons", async () => {
+    query.mockResolvedValueOnce({ rows: [{ count: 0 }] });
+    const store = new PostgreSQLProjectRoomMessageStore({ pool: mockPool });
+    const after = new Date("2026-01-01T00:00:00.000Z");
+    await expect(store.countAfter({ tenantId: "tenant-1", roomId: "room-1", after })).resolves.toBe(0);
+    expect(query.mock.calls[0][1]).toEqual(["tenant-1", "room-1", after, null]);
+    await expect(store.countAfter({ tenantId: "tenant-1", roomId: "room-1", after: new Date(Number.NaN) }))
+      .rejects.toThrow(RangeError);
+    expect(query).toHaveBeenCalledTimes(1);
+  });
+});
+
+const databaseUrl = process.env.TEST_DATABASE_URL;
+const describeIfDatabase = databaseUrl ? describe : describe.skip;
+
+describeIfDatabase("PostgreSQLProjectRoomReadStateStore", () => {
+  let pool: Pool;
+  let store: PostgreSQLProjectRoomReadStateStore;
+  let messages: PostgreSQLProjectRoomMessageStore;
+  const tenantId = `read-state-${randomUUID()}`;
+
+  beforeAll(async () => {
+    pool = new Pool({ connectionString: databaseUrl });
+    const mm = new MigrationManager(pool);
+    mm.register(createProjectRoomTables);
+    mm.register(createProjectRoomReadStates);
+    await mm.migrate();
+    store = new PostgreSQLProjectRoomReadStateStore({ pool });
+    messages = new PostgreSQLProjectRoomMessageStore({ pool });
+  });
+
+  afterAll(async () => {
+    await pool.query("DELETE FROM lattice_project_room_messages WHERE tenant_id = $1", [tenantId]);
+    await pool.query("DELETE FROM lattice_project_room_read_states WHERE tenant_id = $1", [tenantId]);
+    await pool.end();
+  });
+
+  it("returns null before any report and upserts monotonically", async () => {
+    await expect(store.get(tenantId, "room-1", "u1")).resolves.toBeNull();
+    await store.markRead({ tenantId, roomId: "room-1", userId: "u1", lastReadAt: new Date("2026-01-02T00:00:00Z") });
+    await store.markRead({ tenantId, roomId: "room-1", userId: "u1", lastReadAt: new Date("2026-01-01T00:00:00Z") });
+    const state = await store.get(tenantId, "room-1", "u1");
+    expect(state?.lastReadAt.toISOString()).toBe("2026-01-02T00:00:00.000Z");
+  });
+
+  it("counts messages after a horizon excluding one human author", async () => {
+    const roomId = `room-${randomUUID()}`;
+    const base = (id: string, author: ProjectRoomMessage["author"]): Omit<ProjectRoomMessage, "createdAt"> => ({
+      id, tenantId, workspaceId: "w1", projectId: "p1", roomId, author,
+      content: { type: "text", text: id }, mentions: [], source: "user",
+    });
+    await messages.create({ ...base(`m0-${roomId}`, { type: "human", userId: "u2" }) });
+    await new Promise((resolve) => setTimeout(resolve, 20));
+    await messages.create({ ...base(`m1-${roomId}`, { type: "human", userId: "u1" }) });
+    await messages.create({ ...base(`m2-${roomId}`, { type: "human", userId: "u2" }) });
+    await messages.create({ ...base(`m3-${roomId}`, { type: "bot", membershipId: "b1", assistantId: "a1" }) });
+    const first = await messages.findById(tenantId, `m0-${roomId}`);
+    const count = await messages.countAfter({
+      tenantId, roomId, after: first!.createdAt, excludeAuthorUserId: "u1",
+    });
+    expect(count).toBe(2);
+  });
+});

package/src/createPgStoreConfig.ts

@@ -82,12 +82,13 @@ import {
 import { createAgentWebAppsTable } from "./migrations/agent_web_apps_migration";
 import { addAgentWebAppsIdentity } from "./migrations/add_agent_web_apps_identity_migration";
 import { createCapabilityBundlesTable } from "./migrations/capability_bundle_migration";
-import { createProjectRoomTables } from "./migrations/project_room_migration";
+import { createProjectRoomReadStates, createProjectRoomTables } from "./migrations/project_room_migration";
 import { PostgreSQLCapabilityBundleStore } from "./stores/PostgreSQLCapabilityBundleStore";
 import { PostgreSQLProjectRoomStore } from "./stores/PostgreSQLProjectRoomStore";
 import { PostgreSQLProjectMembershipStore } from "./stores/PostgreSQLProjectMembershipStore";
 import { PostgreSQLProjectBotMembershipStore } from "./stores/PostgreSQLProjectBotMembershipStore";
 import { PostgreSQLProjectRoomMessageStore } from "./stores/PostgreSQLProjectRoomMessageStore";
+import { PostgreSQLProjectRoomReadStateStore } from "./stores/PostgreSQLProjectRoomReadStateStore";
 
 export async function createPgStoreConfig(connectionString: string) {
   const pool = new Pool({ connectionString });
@@ -155,6 +156,7 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(addTrustedRunContextColumn);      // v173
   mm.register(addProjectLifecycleEventIndex);   // v175
   mm.register(addAgentWebAppsIdentity);         // v176
+  mm.register(createProjectRoomReadStates);     // v177
 
   await mm.migrate();
 
@@ -197,6 +199,7 @@ export async function createPgStoreConfig(connectionString: string) {
     projectMembership: new PostgreSQLProjectMembershipStore(opts),
     projectBotMembership: new PostgreSQLProjectBotMembershipStore(opts),
     projectRoomMessage: new PostgreSQLProjectRoomMessageStore(opts),
+    projectRoomReadState: new PostgreSQLProjectRoomReadStateStore(opts),
     vectorStoreProvider: new PGVectorStoreProvider(pool, connectionString),
     checkpoint,
   };

package/src/index.ts

@@ -75,6 +75,7 @@ export * from "./stores/PostgreSQLProjectRoomStore";
 export * from "./stores/PostgreSQLProjectMembershipStore";
 export * from "./stores/PostgreSQLProjectBotMembershipStore";
 export * from "./stores/PostgreSQLProjectRoomMessageStore";
+export * from "./stores/PostgreSQLProjectRoomReadStateStore";
 export * from "./migrations/project_room_migration";
 
 export * from "./PGVectorStoreProvider";

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/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/PostgreSQLProjectRoomMessageStore.ts

@@ -170,6 +170,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]);
+  }
+}