@axiom-lattice/pg-stores 4.0.4 → 4.1.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "4.0.4",
+  "version": "4.1.0",
   "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": "10.0.4",
-    "@axiom-lattice/protocols": "6.2.0"
+    "@axiom-lattice/core": "10.1.0",
+    "@axiom-lattice/protocols": "6.3.0"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",

package/src/__tests__/PostgresSharedResourceStore.kind.test.ts

@@ -0,0 +1,210 @@
+import { randomUUID } from "crypto";
+import { Pool } from "pg";
+import type { CreateShareRecord } from "@axiom-lattice/protocols";
+import { PostgresSharedResourceStore } from "../stores/PostgresSharedResourceStore";
+
+const url = process.env.TEST_DATABASE_URL;
+(url ? describe : describe.skip)("PostgresSharedResourceStore kind", () => {
+  const runId = randomUUID();
+  const tenantId = `t_${runId}`;
+  const userId = `u_${runId}`;
+  const workspaceId = `w_${runId}`;
+  const projectId = `p_${runId}`;
+  let store: PostgresSharedResourceStore;
+  let admin: Pool;
+  let tokenCounter = 0;
+
+  const nextToken = (prefix: string): string => {
+    tokenCounter += 1;
+    return `${prefix}_${runId}_${tokenCounter}`;
+  };
+
+  const baseRecord = (token: string) => ({
+    token,
+    title: "Share",
+    createdBy: userId,
+    tenantId,
+    workspaceId,
+    projectId,
+    assistantId: null,
+    vmIsolation: "project",
+    visibility: "public" as const,
+    passwordHash: null,
+    expiresAt: null,
+    maxAccess: null,
+    revoked: false,
+  });
+
+  const mcpAppAddress = {
+    kind: "mcpApp" as const,
+    app: {
+      kind: "plugin" as const,
+      pluginType: "demo",
+      resource: "ui/app.html",
+      toolName: "render",
+      scope: "signed",
+    },
+  };
+
+  beforeAll(async () => {
+    store = new PostgresSharedResourceStore({
+      poolConfig: { connectionString: url! },
+    });
+    admin = new Pool({ connectionString: url });
+    await store.initialize();
+  });
+
+  afterAll(async () => {
+    if (admin) {
+      await admin.query(
+        "DELETE FROM lattice_shared_resources WHERE tenant_id = $1",
+        [tenantId],
+      );
+      await admin.end();
+    }
+    await store.dispose();
+  });
+
+  it("round-trips file and mcpApp records", async () => {
+    const fileToken = nextToken("file");
+    const fileRecord = await store.create({
+      ...baseRecord(fileToken),
+      kind: "file",
+      address: { volume: "project", resourcePath: "docs/readme.md" },
+    } satisfies CreateShareRecord);
+
+    expect(fileRecord.kind).toBe("file");
+    expect(fileRecord.address).toEqual({
+      volume: "project",
+      resourcePath: "docs/readme.md",
+    });
+
+    const foundFile = await store.findByToken(fileToken);
+    expect(foundFile?.kind).toBe("file");
+    if (foundFile?.kind !== "file") {
+      throw new Error("expected file share record");
+    }
+    expect(foundFile.address).toEqual({
+      volume: "project",
+      resourcePath: "docs/readme.md",
+    });
+
+    const mcpToken = nextToken("mcp");
+    const mcpRecord = await store.create({
+      ...baseRecord(mcpToken),
+      kind: "mcpApp",
+      address: mcpAppAddress,
+    } satisfies CreateShareRecord);
+
+    expect(mcpRecord.kind).toBe("mcpApp");
+    if (mcpRecord.kind !== "mcpApp") {
+      throw new Error("expected mcpApp share record");
+    }
+    expect(mcpRecord.address.app.pluginType).toBe("demo");
+
+    const foundMcp = await store.findByToken(mcpToken);
+    expect(foundMcp?.kind).toBe("mcpApp");
+    if (foundMcp?.kind !== "mcpApp") {
+      throw new Error("expected mcpApp share record");
+    }
+    expect(foundMcp.address.app.resource).toBe("ui/app.html");
+  });
+
+  it("lists both kinds for the same user", async () => {
+    const fileToken = nextToken("list_file");
+    await store.create({
+      ...baseRecord(fileToken),
+      kind: "file",
+      address: { volume: "project", resourcePath: "docs/list.md" },
+    } satisfies CreateShareRecord);
+
+    const mcpToken = nextToken("list_mcp");
+    await store.create({
+      ...baseRecord(mcpToken),
+      kind: "mcpApp",
+      address: mcpAppAddress,
+    } satisfies CreateShareRecord);
+
+    const records = await store.listByUser(tenantId, userId);
+    const tokens = records.map((record) => record.token);
+    expect(tokens).toContain(fileToken);
+    expect(tokens).toContain(mcpToken);
+    expect(new Set(records.map((record) => record.kind))).toEqual(
+      new Set(["file", "mcpApp"]),
+    );
+  });
+
+  it("rejects an unknown kind at the database level", async () => {
+    await expect(
+      admin.query(
+        `INSERT INTO lattice_shared_resources
+         (token, kind, address, created_by, tenant_id, workspace_id, project_id)
+         VALUES ($1, $2, $3, $4, $5, $6, $7)`,
+        [
+          nextToken("bogus"),
+          "bogus",
+          JSON.stringify({ volume: "project", resourcePath: "x.txt" }),
+          userId,
+          tenantId,
+          workspaceId,
+          projectId,
+        ],
+      ),
+    ).rejects.toThrow(/lattice_shared_resources_kind_check/);
+  });
+
+  it("rejects mcpApp shares with internal visibility", async () => {
+    await expect(
+      admin.query(
+        `INSERT INTO lattice_shared_resources
+         (token, kind, address, created_by, tenant_id, workspace_id, project_id, visibility)
+         VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`,
+        [
+          nextToken("internal"),
+          "mcpApp",
+          JSON.stringify(mcpAppAddress),
+          userId,
+          tenantId,
+          workspaceId,
+          projectId,
+          "internal",
+        ],
+      ),
+    ).rejects.toThrow(
+      /lattice_shared_resources_mcp_app_visibility_check/,
+    );
+  });
+
+  it("maps a legacy row without kind to file via the column default", async () => {
+    const legacyToken = nextToken("legacy");
+    await admin.query(
+      `INSERT INTO lattice_shared_resources
+       (token, address, created_by, tenant_id, workspace_id, project_id)
+       VALUES ($1, $2, $3, $4, $5, $6)`,
+      [
+        legacyToken,
+        JSON.stringify({ volume: "project", resourcePath: "legacy.txt" }),
+        userId,
+        tenantId,
+        workspaceId,
+        projectId,
+      ],
+    );
+
+    const { rows } = await admin.query<{ kind: string }>(
+      "SELECT kind FROM lattice_shared_resources WHERE token = $1",
+      [legacyToken],
+    );
+    expect(rows[0]?.kind).toBe("file");
+
+    const legacy = await store.findByToken(legacyToken);
+    expect(legacy?.kind).toBe("file");
+    if (legacy?.kind !== "file") {
+      throw new Error("expected legacy file share record");
+    }
+    expect(legacy.address).toEqual({
+      volume: "project",
+      resourcePath: "legacy.txt",
+    });
+  });
+});

package/src/createPgStoreConfig.ts

@@ -66,6 +66,7 @@ import { addA2AKeyAssistantIds } from "./migrations/add_a2a_key_assistant_ids";
 import { addApiKeyGrants } from "./migrations/add_api_key_grants";
 import { createOpenAuditTable } from "./migrations/create_open_audit_table";
 import { addShareOpenGrants } from "./migrations/add_share_open_grants";
+import { addShareKind } from "./migrations/add_share_kind";
 import { taskMigrations } from "./migrations/task_migration";
 import { createMenuItemsTable } from "./migrations/menu_items_migration";
 import { addFileContentType } from "./migrations/menu_items_add_file_type";
@@ -224,6 +225,7 @@ export async function createPgStoreConfig(connectionString: string) {
   mm.register(createAgentRunProjectTaskMigration); // v203
   mm.register(agentRunProjectTaskRunningIndexMigration); // v204
   mm.register(addModelProviderKindDefault);     // v207
+  mm.register(addShareKind);                    // v208
 
   await mm.migrate();
 

package/src/index.ts

@@ -50,6 +50,7 @@ export * from "./migrations/channel_installation_migrations";
 export * from "./migrations/channel_installations_alter_migration";
 export * from "./migrations/channel_bindings_migration";
 export * from "./migrations/shared_resources_migration";
+export * from "./migrations/add_share_kind";
 export * from "./migrations/a2a_api_key_migration";
 export * from "./migrations/agent_web_apps_migration";
 export * from "./migrations/workflow_tracking_migrations";

package/src/migrations/add_share_kind.ts

@@ -0,0 +1,45 @@
+import type { PoolClient } from "pg";
+import type { Migration } from "./migration";
+
+export const addShareKind: Migration = {
+  version: 208,
+  name: "add_share_kind",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_shared_resources
+        ADD COLUMN IF NOT EXISTS kind VARCHAR(16) NOT NULL DEFAULT 'file'
+    `);
+    await client.query(`
+      ALTER TABLE lattice_shared_resources
+        DROP CONSTRAINT IF EXISTS lattice_shared_resources_kind_check
+    `);
+    await client.query(`
+      ALTER TABLE lattice_shared_resources
+        ADD CONSTRAINT lattice_shared_resources_kind_check
+        CHECK (kind IN ('file', 'mcpApp'))
+    `);
+    await client.query(`
+      ALTER TABLE lattice_shared_resources
+        DROP CONSTRAINT IF EXISTS lattice_shared_resources_mcp_app_visibility_check
+    `);
+    await client.query(`
+      ALTER TABLE lattice_shared_resources
+        ADD CONSTRAINT lattice_shared_resources_mcp_app_visibility_check
+        CHECK (kind <> 'mcpApp' OR visibility IN ('public', 'password'))
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_shared_resources
+        DROP CONSTRAINT IF EXISTS lattice_shared_resources_mcp_app_visibility_check
+    `);
+    await client.query(`
+      ALTER TABLE lattice_shared_resources
+        DROP CONSTRAINT IF EXISTS lattice_shared_resources_kind_check
+    `);
+    await client.query(`
+      ALTER TABLE lattice_shared_resources
+        DROP COLUMN IF EXISTS kind
+    `);
+  },
+};

package/src/stores/PostgresSharedResourceStore.ts

@@ -1,9 +1,16 @@
 import { Pool } from "pg";
 import type { PoolConfig } from "pg";
-import type { ShareRecord, ResourceAddress } from "@axiom-lattice/protocols";
+import type {
+  CreateShareRecord,
+  McpAppShareApp,
+  ShareRecord,
+  ShareVisibility,
+} from "@axiom-lattice/protocols";
 import type { SharedResourceStore } from "./SharedResourceStore";
 import { MigrationManager } from "../migrations/migration";
 import { createSharedResourcesTable } from "../migrations/shared_resources_migration";
+import { addShareOpenGrants } from "../migrations/add_share_open_grants";
+import { addShareKind } from "../migrations/add_share_kind";
 
 export interface PostgresSharedResourceStoreOptions {
   pool?: Pool;
@@ -11,9 +18,64 @@ export interface PostgresSharedResourceStoreOptions {
   autoMigrate?: boolean;
 }
 
+/**
+ * Raised when a persisted MCP App share row does not carry a valid app target.
+ */
+export class InvalidMcpAppShareAddressError extends Error {
+  /** Machine-readable failure code. */
+  readonly code = "INVALID_MCP_APP_SHARE_ADDRESS" as const;
+
+  /**
+   * @param token - Share token owning the invalid address.
+   */
+  constructor(token: string) {
+    super(`Share ${token} has an invalid MCP App address`);
+    this.name = "InvalidMcpAppShareAddressError";
+  }
+}
+
+/**
+ * Raised when a persisted MCP App share row uses an unsupported visibility.
+ */
+export class InvalidMcpAppShareVisibilityError extends Error {
+  /** Machine-readable failure code. */
+  readonly code = "INVALID_MCP_APP_SHARE_VISIBILITY" as const;
+  /** Visibility value found in the row. */
+  readonly visibility: string;
+
+  /**
+   * @param token - Share token owning the invalid visibility.
+   * @param visibility - Visibility value found in the row.
+   */
+  constructor(token: string, visibility: string) {
+    super(
+      `Share ${token} has unsupported MCP App visibility: ${visibility}`,
+    );
+    this.name = "InvalidMcpAppShareVisibilityError";
+    this.visibility = visibility;
+  }
+}
+
+/**
+ * Raised when a persisted file share row does not carry a valid resource address.
+ */
+export class InvalidFileShareAddressError extends Error {
+  /** Machine-readable failure code. */
+  readonly code = "INVALID_FILE_SHARE_ADDRESS" as const;
+
+  /**
+   * @param token - Share token owning the invalid address.
+   */
+  constructor(token: string) {
+    super(`Share ${token} has an invalid file address`);
+    this.name = "InvalidFileShareAddressError";
+  }
+}
+
 type ShareRow = {
   token: string;
-  address: ResourceAddress;
+  kind: string | null;
+  address: unknown;
   title: string | null;
   created_by: string;
   tenant_id: string;
@@ -57,6 +119,8 @@ export class PostgresSharedResourceStore implements SharedResourceStore {
 
     this.migrationManager = new MigrationManager(this.pool);
     this.migrationManager.register(createSharedResourcesTable);
+    this.migrationManager.register(addShareOpenGrants);
+    this.migrationManager.register(addShareKind);
 
     if (options.autoMigrate !== false) {
       this.initialize().catch((error) => {
@@ -99,19 +163,18 @@ export class PostgresSharedResourceStore implements SharedResourceStore {
     return this._mapRow(rows[0]);
   }
 
-  async create(
-    record: Omit<ShareRecord, "accessCount" | "createdAt" | "version">,
-  ): Promise<ShareRecord> {
+  async create(record: CreateShareRecord): Promise<ShareRecord> {
     await this.ensureInitialized();
     const { rows } = await this.pool.query<ShareRow>(
       `INSERT INTO lattice_shared_resources
-       (token, address, title, created_by, tenant_id, workspace_id, project_id,
+       (token, kind, address, title, created_by, tenant_id, workspace_id, project_id,
         assistant_id, vm_isolation, visibility, password_hash,
         expires_at, max_access, open_grants, revoked)
-       VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15)
+       VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16)
        RETURNING *`,
       [
         record.token,
+        record.kind ?? "file",
         JSON.stringify(record.address),
         record.title,
         record.createdBy,
@@ -210,16 +273,9 @@ export class PostgresSharedResourceStore implements SharedResourceStore {
   }
 
   private _mapRow(row: ShareRow): ShareRecord {
-    let address: ResourceAddress;
-    if (typeof row.address === "string") {
-      address = JSON.parse(row.address) as ResourceAddress;
-    } else {
-      address = row.address as ResourceAddress;
-    }
-
-    return {
+    const address = parseShareAddress(row.address);
+    const base = {
       token: row.token,
-      address,
       title: row.title,
       createdBy: row.created_by,
       tenantId: row.tenant_id,
@@ -227,7 +283,6 @@ export class PostgresSharedResourceStore implements SharedResourceStore {
       projectId: row.project_id,
       assistantId: row.assistant_id,
       vmIsolation: row.vm_isolation,
-      visibility: row.visibility as ShareRecord["visibility"],
       passwordHash: row.password_hash,
       expiresAt: row.expires_at ? new Date(row.expires_at) : null,
       maxAccess: row.max_access,
@@ -237,5 +292,58 @@ export class PostgresSharedResourceStore implements SharedResourceStore {
       createdAt: new Date(row.created_at),
       version: row.version ?? 1,
     };
+
+    if (row.kind === "mcpApp") {
+      if (!isRecord(address) || address.kind !== "mcpApp") {
+        throw new InvalidMcpAppShareAddressError(row.token);
+      }
+      const app = address.app;
+      if (!isMcpAppShareApp(app)) {
+        throw new InvalidMcpAppShareAddressError(row.token);
+      }
+      if (row.visibility !== "public" && row.visibility !== "password") {
+        throw new InvalidMcpAppShareVisibilityError(row.token, row.visibility);
+      }
+      return {
+        ...base,
+        kind: "mcpApp",
+        address: { kind: "mcpApp", app },
+        visibility: row.visibility,
+      };
+    }
+
+    if (
+      !isRecord(address) ||
+      typeof address.volume !== "string" ||
+      typeof address.resourcePath !== "string"
+    ) {
+      throw new InvalidFileShareAddressError(row.token);
+    }
+    return {
+      ...base,
+      kind: "file",
+      address: {
+        volume: address.volume,
+        resourcePath: address.resourcePath,
+      },
+      visibility: row.visibility as ShareVisibility,
+    };
+  }
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function isMcpAppShareApp(value: unknown): value is McpAppShareApp {
+  return isRecord(value) && (value.kind === "plugin" || value.kind === "mcp");
+}
+
+function parseShareAddress(raw: unknown): unknown {
+  if (typeof raw !== "string") return raw;
+  try {
+    return JSON.parse(raw) as unknown;
+  } catch {
+    return null;
   }
 }