@axiom-lattice/pg-stores 3.3.12 → 3.4.0

package/src/stores/PostgreSQLStore.ts

@@ -0,0 +1,167 @@
+import type { Pool } from "pg";
+import {
+  assertStoreIdentity,
+  assertStoreNamespacePrefix,
+  storeExpiresAt,
+  storePaginationValue,
+  type StoreItem,
+  type StoreProtocol,
+  type StorePutInput,
+  type StoreSearchInput,
+} from "@axiom-lattice/protocols";
+
+interface StoreRow {
+  tenant_id: string;
+  namespace: string[];
+  key: string;
+  value: Record<string, unknown> | null;
+  created_at: Date;
+  updated_at: Date;
+  expires_at: Date | null;
+}
+
+function mapRow(row: StoreRow): StoreItem {
+  return {
+    namespace: row.namespace,
+    key: row.key,
+    value: row.value ?? {},
+    createdAt: row.created_at,
+    updatedAt: row.updated_at,
+    expiresAt: row.expires_at ?? undefined,
+  };
+}
+
+/**
+ * PostgreSQL tenant-scoped namespaced key-value store.
+ *
+ * @remarks
+ * - `namespace` is a `TEXT[]`; prefix search compares `namespace[1:n]`.
+ * - `filter` uses JSONB containment (`value @> $filter`).
+ * - `query` is a case-insensitive substring over `value::text`.
+ * - Expired items are excluded from reads/searches and hard-deleted by
+ *   {@link deleteExpired}.
+ */
+export class PostgreSQLStore implements StoreProtocol {
+  /**
+   * @param pool - Pool shared with the rest of the store config.
+   */
+  constructor(private readonly pool: Pool) {}
+
+  async put(tenantId: string, input: StorePutInput): Promise<StoreItem> {
+    assertStoreIdentity(input.namespace, input.key);
+    const now = new Date();
+    const expiresAt = storeExpiresAt(input.ttlMinutes, now) ?? null;
+    const result = await this.pool.query<StoreRow>(
+      `INSERT INTO lattice_store_items
+        (tenant_id, namespace, key, value, created_at, updated_at, expires_at)
+       VALUES ($1, $2, $3, $4::jsonb, $5, $5, $6)
+       ON CONFLICT (tenant_id, namespace, key) DO UPDATE SET
+         value = EXCLUDED.value,
+         updated_at = EXCLUDED.updated_at,
+         expires_at = EXCLUDED.expires_at
+       RETURNING tenant_id, namespace, key, value, created_at, updated_at, expires_at`,
+      [tenantId, input.namespace, input.key, JSON.stringify(input.value), now, expiresAt],
+    );
+    return mapRow(result.rows[0]);
+  }
+
+  async get(tenantId: string, namespace: string[], key: string): Promise<StoreItem | null> {
+    assertStoreIdentity(namespace, key);
+    const result = await this.pool.query<StoreRow>(
+      `SELECT tenant_id, namespace, key, value, created_at, updated_at, expires_at
+       FROM lattice_store_items
+       WHERE tenant_id = $1 AND namespace = $2 AND key = $3
+         AND (expires_at IS NULL OR expires_at > $4)`,
+      [tenantId, namespace, key, new Date()],
+    );
+    return result.rows[0] ? mapRow(result.rows[0]) : null;
+  }
+
+  async delete(tenantId: string, namespace: string[], key: string): Promise<boolean> {
+    assertStoreIdentity(namespace, key);
+    const result = await this.pool.query(
+      `DELETE FROM lattice_store_items
+       WHERE tenant_id = $1 AND namespace = $2 AND key = $3`,
+      [tenantId, namespace, key],
+    );
+    return (result.rowCount ?? 0) > 0;
+  }
+
+  async search(tenantId: string, input: StoreSearchInput): Promise<StoreItem[]> {
+    const conditions = [
+      "tenant_id = $1",
+      "(expires_at IS NULL OR expires_at > $2)",
+    ];
+    const params: unknown[] = [tenantId, new Date()];
+    if (input.namespacePrefix) {
+      assertStoreNamespacePrefix(input.namespacePrefix);
+    }
+    if (input.namespacePrefix && input.namespacePrefix.length > 0) {
+      params.push(input.namespacePrefix);
+      conditions.push(`namespace[1:array_length($${params.length}::text[], 1)] = $${params.length}`);
+    }
+    if (input.filter !== undefined) {
+      params.push(JSON.stringify(input.filter));
+      conditions.push(`value @> $${params.length}::jsonb`);
+    }
+    if (input.query !== undefined && input.query !== "") {
+      params.push(input.query.replace(/[\\%_]/g, "\\$&"));
+      conditions.push(`value::text ILIKE '%' || $${params.length} || '%' ESCAPE '\\'`);
+    }
+    params.push(storePaginationValue(input.limit, 10));
+    const limitParam = params.length;
+    params.push(storePaginationValue(input.offset, 0));
+    const offsetParam = params.length;
+    const result = await this.pool.query<StoreRow>(
+      `SELECT tenant_id, namespace, key, value, created_at, updated_at, expires_at
+       FROM lattice_store_items
+       WHERE ${conditions.join(" AND ")}
+       ORDER BY updated_at DESC, namespace, key
+       LIMIT $${limitParam} OFFSET $${offsetParam}`,
+      params,
+    );
+    return result.rows.map(mapRow);
+  }
+
+  async listNamespaces(
+    tenantId: string,
+    prefix: string[] = [],
+    limit = 100,
+  ): Promise<string[][]> {
+    assertStoreNamespacePrefix(prefix);
+    const conditions = [
+      "tenant_id = $1",
+      "(expires_at IS NULL OR expires_at > $2)",
+    ];
+    const params: unknown[] = [tenantId, new Date()];
+    if (prefix.length > 0) {
+      params.push(prefix);
+      conditions.push(`namespace[1:array_length($${params.length}::text[], 1)] = $${params.length}`);
+    }
+    params.push(storePaginationValue(limit, 100));
+    const result = await this.pool.query<{ namespace: string[] }>(
+      `SELECT DISTINCT namespace
+       FROM lattice_store_items
+       WHERE ${conditions.join(" AND ")}
+       ORDER BY namespace
+       LIMIT $${params.length}`,
+      params,
+    );
+    return result.rows.map((row) => row.namespace);
+  }
+
+  async deleteExpired(now: Date, limit: number): Promise<number> {
+    const bounded = storePaginationValue(limit, 0);
+    const result = await this.pool.query(
+      `WITH expired AS (
+         SELECT ctid FROM lattice_store_items
+         WHERE expires_at IS NOT NULL AND expires_at <= $1
+         ORDER BY expires_at
+         LIMIT $2
+       )
+       DELETE FROM lattice_store_items WHERE ctid IN (SELECT ctid FROM expired)`,
+      [now, bounded],
+    );
+    return result.rowCount ?? 0;
+  }
+}

package/src/stores/PostgreSQLThreadStore.ts

@@ -10,7 +10,7 @@ import {
   CreateThreadRequest,
 } from "@axiom-lattice/protocols";
 import { MigrationManager } from "../migrations/migration";
-import { createThreadsTable } from "../migrations/thread_migrations";
+import { addThreadExpiresAt, createThreadsTable } from "../migrations/thread_migrations";
 import { addThreadTenantId } from "../migrations/thread_tenant_migration";
 import { changeThreadPrimaryKey } from "../migrations/thread_pk_migration";
 
@@ -72,6 +72,7 @@ export class PostgreSQLThreadStore implements ThreadStore {
     this.migrationManager.register(createThreadsTable);
     this.migrationManager.register(addThreadTenantId);
     this.migrationManager.register(changeThreadPrimaryKey);
+    this.migrationManager.register(addThreadExpiresAt);
 
     // Auto-migrate by default
     if (options.autoMigrate !== false) {
@@ -134,9 +135,10 @@ export class PostgreSQLThreadStore implements ThreadStore {
       metadata: any;
       created_at: Date;
       updated_at: Date;
+      expires_at: Date | null;
     }>(
       `
-      SELECT id, tenant_id, assistant_id, metadata, created_at, updated_at
+      SELECT id, tenant_id, assistant_id, metadata, created_at, updated_at, expires_at
       FROM lattice_threads
       WHERE tenant_id = $1 AND assistant_id = $2
       ORDER BY created_at DESC
@@ -174,9 +176,10 @@ export class PostgreSQLThreadStore implements ThreadStore {
       metadata: any;
       created_at: Date;
       updated_at: Date;
+      expires_at: Date | null;
     }>(
       `
-      SELECT id, tenant_id, assistant_id, metadata, created_at, updated_at
+      SELECT id, tenant_id, assistant_id, metadata, created_at, updated_at, expires_at
       FROM lattice_threads
       WHERE tenant_id = $1 AND id = $2
     `,
@@ -206,14 +209,15 @@ export class PostgreSQLThreadStore implements ThreadStore {
 
     await this.pool.query(
       `
-      INSERT INTO lattice_threads (id, tenant_id, assistant_id, metadata, created_at, updated_at)
-      VALUES ($1, $2, $3, $4, $5, $6)
+      INSERT INTO lattice_threads (id, tenant_id, assistant_id, metadata, created_at, updated_at, expires_at)
+      VALUES ($1, $2, $3, $4, $5, $6, $7)
       ON CONFLICT (id, tenant_id) DO UPDATE SET
         assistant_id = EXCLUDED.assistant_id,
         metadata = EXCLUDED.metadata,
-        updated_at = EXCLUDED.updated_at
+        updated_at = EXCLUDED.updated_at,
+        expires_at = EXCLUDED.expires_at
     `,
-      [threadId, tenantId, assistantId, JSON.stringify(metadata), now, now]
+      [threadId, tenantId, assistantId, JSON.stringify(metadata), now, now, data.expiresAt ?? null]
     );
 
     return {
@@ -223,6 +227,7 @@ export class PostgreSQLThreadStore implements ThreadStore {
       metadata,
       createdAt: now,
       updatedAt: now,
+      expiresAt: data.expiresAt,
     };
   }
 
@@ -249,20 +254,24 @@ export class PostgreSQLThreadStore implements ThreadStore {
     };
 
     const now = new Date();
+    const expiresAt = updates.expiresAt !== undefined
+      ? updates.expiresAt
+      : existing.expiresAt ?? null;
 
     await this.pool.query(
       `
       UPDATE lattice_threads
-      SET metadata = $1, updated_at = $2
-      WHERE tenant_id = $3 AND id = $4
+      SET metadata = $1, updated_at = $2, expires_at = $3
+      WHERE tenant_id = $4 AND id = $5
     `,
-      [JSON.stringify(updatedMetadata), now, tenantId, threadId]
+      [JSON.stringify(updatedMetadata), now, expiresAt, tenantId, threadId]
     );
 
     return {
       ...existing,
       metadata: updatedMetadata,
       updatedAt: now,
+      expiresAt: expiresAt ?? undefined,
     };
   }
 
@@ -301,6 +310,34 @@ export class PostgreSQLThreadStore implements ThreadStore {
     return result.rows.length > 0;
   }
 
+  /**
+   * List expired threads across all tenants, oldest expiry first.
+   */
+  async listExpiredThreads(now: Date, limit: number): 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;
+      expires_at: Date | null;
+    }>(
+      `
+      SELECT id, tenant_id, assistant_id, metadata, created_at, updated_at, expires_at
+      FROM lattice_threads
+      WHERE expires_at IS NOT NULL AND expires_at <= $1
+      ORDER BY expires_at ASC
+      LIMIT $2
+    `,
+      [now, Math.max(0, limit)]
+    );
+
+    return result.rows.map(this.mapRowToThread);
+  }
+
   /**
    * Ensure store is initialized
    */
@@ -320,6 +357,7 @@ export class PostgreSQLThreadStore implements ThreadStore {
     metadata: any;
     created_at: Date;
     updated_at: Date;
+    expires_at: Date | null;
   }): Thread {
     return {
       id: row.id,
@@ -331,6 +369,7 @@ export class PostgreSQLThreadStore implements ThreadStore {
           : row.metadata || {},
       createdAt: row.created_at,
       updatedAt: row.updated_at,
+      expiresAt: row.expires_at ?? undefined,
     };
   }
 }

package/src/stores/ThreadMessageQueueStore.ts

@@ -277,21 +277,26 @@ export class ThreadMessageQueueStore implements IMessageQueueStore {
 
   /**
    * Get all threads with pending or processing messages
+   *
+   * @returns One entry per full scope, with `hasPending` distinguishing
+   * processing-only threads (no pending row) from threads with work to drain.
    */
   async getThreadsWithPendingMessages(): Promise<ThreadInfo[]> {
     const result = await this.pool.query(
-       `SELECT DISTINCT ON (tenant_id, assistant_id, workspace_id, project_id, thread_id) tenant_id, assistant_id, thread_id, workspace_id, project_id
-       FROM lattice_thread_message_queue 
-       WHERE status IN ('pending', 'processing')
-        ORDER BY tenant_id, assistant_id, workspace_id, project_id, thread_id`
+      `SELECT tenant_id, assistant_id, thread_id, workspace_id, project_id,
+              bool_or(status = 'pending') AS has_pending
+         FROM lattice_thread_message_queue
+        WHERE status IN ('pending', 'processing')
+        GROUP BY tenant_id, assistant_id, thread_id, workspace_id, project_id`
     );
 
-    return result.rows.map(row => ({
+    return result.rows.map((row) => ({
       tenantId: row.tenant_id,
       assistantId: row.assistant_id,
       threadId: row.thread_id,
       workspaceId: row.workspace_id,
       projectId: row.project_id,
+      hasPending: row.has_pending,
     }));
   }