@axiom-lattice/pg-stores 1.0.21 → 1.0.23

package/src/stores/PostgreSQLThreadStore.ts

@@ -10,6 +10,8 @@ import {
 } from "@axiom-lattice/protocols";
 import { MigrationManager } from "../migrations/migration";
 import { createThreadsTable } from "../migrations/thread_migrations";
+import { addThreadTenantId } from "../migrations/thread_tenant_migration";
+import { changeThreadPrimaryKey } from "../migrations/thread_pk_migration";
 
 /**
  * PostgreSQL ThreadStore options
@@ -30,13 +32,16 @@ export interface PostgreSQLThreadStoreOptions {
 
 /**
  * PostgreSQL implementation of ThreadStore
+ * 
+ * Features:
+ * - Multi-tenant isolation via tenant_id
  */
 export class PostgreSQLThreadStore implements ThreadStore {
   private pool: Pool;
   private migrationManager: MigrationManager;
   private initialized: boolean = false;
   private ownsPool: boolean = true;
-  private initPromise: Promise<void> | null = null; // Promise-based lock to prevent concurrent initialization
+  private initPromise: Promise<void> | null = null;
 
   constructor(options: PostgreSQLThreadStoreOptions) {
     // Create Pool from config
@@ -48,6 +53,8 @@ export class PostgreSQLThreadStore implements ThreadStore {
 
     this.migrationManager = new MigrationManager(this.pool);
     this.migrationManager.register(createThreadsTable);
+    this.migrationManager.register(addThreadTenantId);
+    this.migrationManager.register(changeThreadPrimaryKey);
 
     // Auto-migrate by default
     if (options.autoMigrate !== false) {
@@ -73,23 +80,19 @@ export class PostgreSQLThreadStore implements ThreadStore {
    * Uses a promise-based lock to prevent concurrent initialization
    */
   async initialize(): Promise<void> {
-    // If already initialized, return immediately
     if (this.initialized) {
       return;
     }
 
-    // If initialization is in progress, wait for it
     if (this.initPromise) {
       return this.initPromise;
     }
 
-    // Start initialization and store the promise
     this.initPromise = (async () => {
       try {
         await this.migrationManager.migrate();
         this.initialized = true;
       } finally {
-        // Clear the promise after completion (success or failure)
         this.initPromise = null;
       }
     })();
@@ -98,52 +101,54 @@ export class PostgreSQLThreadStore implements ThreadStore {
   }
 
   /**
-   * Get all threads for a specific assistant
+   * Get all threads for a specific tenant and assistant
    */
-  async getThreadsByAssistantId(assistantId: string): Promise<Thread[]> {
+  async getThreadsByAssistantId(tenantId: string, assistantId: string): Promise<Thread[]> {
     await this.ensureInitialized();
 
     const result = await this.pool.query<{
       id: string;
+      tenant_id: string;
       assistant_id: string;
       metadata: any;
       created_at: Date;
       updated_at: Date;
     }>(
       `
-      SELECT id, assistant_id, metadata, created_at, updated_at
+      SELECT id, tenant_id, assistant_id, metadata, created_at, updated_at
       FROM lattice_threads
-      WHERE assistant_id = $1
+      WHERE tenant_id = $1 AND assistant_id = $2
       ORDER BY created_at DESC
     `,
-      [assistantId]
+      [tenantId, assistantId]
     );
 
     return result.rows.map(this.mapRowToThread);
   }
 
   /**
-   * Get a thread by ID for a specific assistant
+   * Get a thread by ID for a specific tenant
    */
   async getThreadById(
-    assistantId: string,
+    tenantId: string,
     threadId: string
   ): Promise<Thread | undefined> {
     await this.ensureInitialized();
 
     const result = await this.pool.query<{
       id: string;
+      tenant_id: string;
       assistant_id: string;
       metadata: any;
       created_at: Date;
       updated_at: Date;
     }>(
       `
-      SELECT id, assistant_id, metadata, created_at, updated_at
+      SELECT id, tenant_id, assistant_id, metadata, created_at, updated_at
       FROM lattice_threads
-      WHERE id = $1 AND assistant_id = $2
+      WHERE tenant_id = $1 AND id = $2
     `,
-      [threadId, assistantId]
+      [tenantId, threadId]
     );
 
     if (result.rows.length === 0) {
@@ -154,9 +159,10 @@ export class PostgreSQLThreadStore implements ThreadStore {
   }
 
   /**
-   * Create a new thread for an assistant
+   * Create a new thread for a tenant and assistant
    */
   async createThread(
+    tenantId: string,
     assistantId: string,
     threadId: string,
     data: CreateThreadRequest
@@ -168,17 +174,19 @@ export class PostgreSQLThreadStore implements ThreadStore {
 
     await this.pool.query(
       `
-      INSERT INTO lattice_threads (id, assistant_id, metadata, created_at, updated_at)
-      VALUES ($1, $2, $3, $4, $5)
-      ON CONFLICT (id, assistant_id) DO UPDATE SET
+      INSERT INTO lattice_threads (id, tenant_id, assistant_id, metadata, created_at, updated_at)
+      VALUES ($1, $2, $3, $4, $5, $6)
+      ON CONFLICT (id, tenant_id) DO UPDATE SET
+        assistant_id = EXCLUDED.assistant_id,
         metadata = EXCLUDED.metadata,
         updated_at = EXCLUDED.updated_at
     `,
-      [threadId, assistantId, JSON.stringify(metadata), now, now]
+      [threadId, tenantId, assistantId, JSON.stringify(metadata), now, now]
     );
 
     return {
       id: threadId,
+      tenantId,
       assistantId,
       metadata,
       createdAt: now,
@@ -190,14 +198,14 @@ export class PostgreSQLThreadStore implements ThreadStore {
    * Update an existing thread
    */
   async updateThread(
-    assistantId: string,
+    tenantId: string,
     threadId: string,
     updates: Partial<CreateThreadRequest>
   ): Promise<Thread | null> {
     await this.ensureInitialized();
 
     // Get existing thread
-    const existing = await this.getThreadById(assistantId, threadId);
+    const existing = await this.getThreadById(tenantId, threadId);
     if (!existing) {
       return null;
     }
@@ -214,9 +222,9 @@ export class PostgreSQLThreadStore implements ThreadStore {
       `
       UPDATE lattice_threads
       SET metadata = $1, updated_at = $2
-      WHERE id = $3 AND assistant_id = $4
+      WHERE tenant_id = $3 AND id = $4
     `,
-      [JSON.stringify(updatedMetadata), now, threadId, assistantId]
+      [JSON.stringify(updatedMetadata), now, tenantId, threadId]
     );
 
     return {
@@ -229,15 +237,15 @@ export class PostgreSQLThreadStore implements ThreadStore {
   /**
    * Delete a thread by ID
    */
-  async deleteThread(assistantId: string, threadId: string): Promise<boolean> {
+  async deleteThread(tenantId: string, threadId: string): Promise<boolean> {
     await this.ensureInitialized();
 
     const result = await this.pool.query(
       `
       DELETE FROM lattice_threads
-      WHERE id = $1 AND assistant_id = $2
+      WHERE tenant_id = $1 AND id = $2
     `,
-      [threadId, assistantId]
+      [tenantId, threadId]
     );
 
     return result.rowCount !== null && result.rowCount > 0;
@@ -246,16 +254,16 @@ export class PostgreSQLThreadStore implements ThreadStore {
   /**
    * Check if thread exists
    */
-  async hasThread(assistantId: string, threadId: string): Promise<boolean> {
+  async hasThread(tenantId: string, threadId: string): Promise<boolean> {
     await this.ensureInitialized();
 
     const result = await this.pool.query(
       `
       SELECT 1 FROM lattice_threads
-      WHERE id = $1 AND assistant_id = $2
+      WHERE tenant_id = $1 AND id = $2
       LIMIT 1
     `,
-      [threadId, assistantId]
+      [tenantId, threadId]
     );
 
     return result.rows.length > 0;
@@ -275,6 +283,7 @@ export class PostgreSQLThreadStore implements ThreadStore {
    */
   private mapRowToThread(row: {
     id: string;
+    tenant_id: string;
     assistant_id: string;
     metadata: any;
     created_at: Date;
@@ -282,6 +291,7 @@ export class PostgreSQLThreadStore implements ThreadStore {
   }): Thread {
     return {
       id: row.id,
+      tenantId: row.tenant_id,
       assistantId: row.assistant_id,
       metadata:
         typeof row.metadata === "string"