@axiom-lattice/pg-stores 1.0.70 → 1.0.72

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/pg-stores",
-  "version": "1.0.70",
+  "version": "1.0.72",
   "description": "PG stores implementation for Axiom Lattice framework",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -20,10 +20,11 @@
   "author": "",
   "license": "MIT",
   "dependencies": {
+    "@langchain/langgraph-checkpoint-postgres": "1.0.0",
     "pg": "^8.16.3",
     "uuid": "^9.0.1",
-    "@axiom-lattice/core": "2.1.79",
-    "@axiom-lattice/protocols": "2.1.41"
+    "@axiom-lattice/core": "2.1.81",
+    "@axiom-lattice/protocols": "2.1.43"
   },
   "devDependencies": {
     "@types/node": "^20.11.24",
@@ -36,7 +37,7 @@
   "scripts": {
     "build": "tsup src/index.ts --format cjs,esm --dts --sourcemap",
     "dev": "tsup src/index.ts --format cjs,esm --dts --watch",
-    "lint": "tsc --noEmit",
+    "lint": "tsc --noEmit && node scripts/check-migration-versions.mjs",
     "clean": "rimraf dist"
   }
 }

package/scripts/check-migration-versions.mjs

@@ -0,0 +1,37 @@
+/**
+ * Check for duplicate migration version numbers.
+ * Exits with code 1 if any version collision is found.
+ * Run as part of lint pipeline.
+ */
+import { readFileSync, readdirSync } from "node:fs";
+import { join, dirname } from "node:path";
+import { fileURLToPath } from "node:url";
+
+const MIGRATIONS_DIR = join(dirname(fileURLToPath(import.meta.url)), "../src/migrations");
+
+const files = readdirSync(MIGRATIONS_DIR).filter(f => f.endsWith(".ts"));
+const versionMap = new Map();
+
+let hasError = false;
+
+for (const file of files) {
+  const content = readFileSync(join(MIGRATIONS_DIR, file), "utf-8");
+  const matches = content.matchAll(/^\s*version:\s*(\d+)/gm);
+  for (const m of matches) {
+    const version = parseInt(m[1]);
+    if (versionMap.has(version)) {
+      console.error(
+        `\x1b[31mMIGRATION VERSION CONFLICT:\x1b[0m v${version} in both "${versionMap.get(version)}" and "${file}"`
+      );
+      hasError = true;
+    } else {
+      versionMap.set(version, file);
+    }
+  }
+}
+
+if (hasError) {
+  process.exit(1);
+}
+
+console.log(`\x1b[32mMigration versions OK\x1b[0m (${versionMap.size} unique versions in ${files.length} files)`);

package/src/createPgStoreConfig.ts

@@ -1,10 +1,3 @@
-/**
- * createPgStoreConfig
- *
- * Creates a map of all PG store instances from a single connection string,
- * ready to pass directly to configureStores().
- */
-
 import { PostgreSQLThreadStore } from "./stores/PostgreSQLThreadStore";
 import { PostgreSQLAssistantStore } from "./stores/PostgreSQLAssistantStore";
 import { PostgreSQLDatabaseConfigStore } from "./stores/PostgreSQLDatabaseConfigStore";
@@ -22,28 +15,38 @@ import { ChannelBindingStore } from "./stores/ChannelBindingStore";
 import { PostgreSQLChannelInstallationStore } from "./stores/PostgreSQLChannelInstallationStore";
 import { PostgreSQLA2AApiKeyStore } from "./stores/PostgreSQLA2AApiKeyStore";
 import { PostgreSQLScheduleStorage } from "./stores/PostgreSQLScheduleStorage";
+import { PostgreSQLTaskStore } from "./stores/PostgreSQLTaskStore";
+import { MenuStore } from "./stores/MenuStore";
+import { PostgresSaver } from "@langchain/langgraph-checkpoint-postgres";
 
 export function createPgStoreConfig(connectionString: string) {
-  const opts = { poolConfig: connectionString, autoMigrate: false } as const;
-  const optsAuto = { poolConfig: connectionString, autoMigrate: true } as const;
+  const opts = { poolConfig: connectionString, autoMigrate: true } as const;
+
+  const checkpoint = PostgresSaver.fromConnString(connectionString);
+  checkpoint.setup().catch((err) => {
+    console.error("[pg-stores] Failed to setup checkpoint table:", err.message || err);
+  });
 
   return {
     workspace: new PostgreSQLWorkspaceStore(opts),
     project: new PostgreSQLProjectStore(opts),
-    eval: new PostgreSQLEvalStore(optsAuto),
+    eval: new PostgreSQLEvalStore(opts),
     user: new PostgreSQLUserStore(opts),
     tenant: new PostgreSQLTenantStore(opts),
     userTenantLink: new PostgreSQLUserTenantLinkStore(opts),
-    channelBinding: new ChannelBindingStore(optsAuto),
-    channelInstallation: new PostgreSQLChannelInstallationStore(optsAuto),
+    channelBinding: new ChannelBindingStore(opts),
+    channelInstallation: new PostgreSQLChannelInstallationStore(opts),
     thread: new PostgreSQLThreadStore(opts),
     database: new PostgreSQLDatabaseConfigStore(opts),
     metrics: new PostgreSQLMetricsServerConfigStore(opts),
     mcp: new PostgreSQLMcpServerConfigStore(opts),
     assistant: new PostgreSQLAssistantStore(opts),
-    workflowTracking: new PostgreSQLWorkflowTrackingStore(optsAuto),
-    threadMessageQueue: new ThreadMessageQueueStore(optsAuto),
-    a2aApiKey: new PostgreSQLA2AApiKeyStore(optsAuto),
+    workflowTracking: new PostgreSQLWorkflowTrackingStore(opts),
+    threadMessageQueue: new ThreadMessageQueueStore(opts),
+    task: new PostgreSQLTaskStore(opts),
+    a2aApiKey: new PostgreSQLA2AApiKeyStore(opts),
     schedule: new PostgreSQLScheduleStorage(opts),
+    menu: new MenuStore(opts),
+    checkpoint,
   };
-}
+}

package/src/index.ts

@@ -57,6 +57,8 @@ export * from "./stores/PostgreSQLA2AApiKeyStore";
 export * from "./stores/PostgreSQLWorkflowTrackingStore";
 export * from "./stores/PostgreSQLEvalStore";
 export * from "./migrations/eval_migrations";
+export * from "./migrations/menu_items_migration";
+export * from "./stores/MenuStore";
 
 // Re-export for convenience
 export { PostgreSQLThreadStore } from "./stores/PostgreSQLThreadStore";
@@ -72,6 +74,7 @@ export { PostgreSQLUserStore } from "./stores/PostgreSQLUserStore";
 export { PostgreSQLTenantStore } from "./stores/PostgreSQLTenantStore";
 export { PostgreSQLUserTenantLinkStore } from "./stores/PostgreSQLUserTenantLinkStore";
 export { ChannelBindingStore } from "./stores/ChannelBindingStore";
+export { MenuStore } from "./stores/MenuStore";
 export { ChannelIdentityMappingStore } from "./stores/ChannelIdentityMappingStore";
 export { PostgreSQLChannelInstallationStore } from "./stores/PostgreSQLChannelInstallationStore";
 export { PostgreSQLA2AApiKeyStore } from "./stores/PostgreSQLA2AApiKeyStore";
@@ -102,6 +105,9 @@ export type {
   Binding,
   BindingRegistry,
   CreateBindingInput,
+  MenuItem,
+  MenuRegistry,
+  CreateMenuItemInput,
   ChannelInstallationStore,
   ChannelInstallation,
   CreateChannelInstallationRequest,

package/src/migrations/add_workspace_project_to_queue.ts

@@ -0,0 +1,30 @@
+/**
+ * Add workspace_id and project_id columns to thread message queue
+ * so that thread recovery after server restart uses the correct scope.
+ */
+
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+export const addWorkspaceProjectToQueue: Migration = {
+  version: 131,
+  name: "add_workspace_project_to_queue",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_thread_message_queue
+      ADD COLUMN IF NOT EXISTS workspace_id VARCHAR(255)
+    `);
+    await client.query(`
+      ALTER TABLE lattice_thread_message_queue
+      ADD COLUMN IF NOT EXISTS project_id VARCHAR(255)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query(
+      "ALTER TABLE lattice_thread_message_queue DROP COLUMN IF EXISTS workspace_id"
+    );
+    await client.query(
+      "ALTER TABLE lattice_thread_message_queue DROP COLUMN IF EXISTS project_id"
+    );
+  },
+};

package/src/migrations/assistant_owner_user_migration.ts

@@ -0,0 +1,26 @@
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+export const addAssistantOwnerUserId: Migration = {
+  version: 132,
+  name: "add_assistant_owner_user_id",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      ALTER TABLE lattice_assistants
+      ADD COLUMN IF NOT EXISTS owner_user_id VARCHAR(255)
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_assistants_owner_user_id
+      ON lattice_assistants(owner_user_id)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query(
+      "DROP INDEX IF EXISTS idx_lattice_assistants_owner_user_id"
+    );
+    await client.query(
+      "ALTER TABLE lattice_assistants DROP COLUMN IF EXISTS owner_user_id"
+    );
+  },
+};

package/src/migrations/menu_items_migration.ts

@@ -0,0 +1,38 @@
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+export const createMenuItemsTable: Migration = {
+  version: 134,
+  name: "create_menu_items_table",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_menu_items (
+        id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+        tenant_id VARCHAR(255) NOT NULL,
+        menu_target VARCHAR(20) NOT NULL CHECK (menu_target IN ('sidebar', 'workspace')),
+        "group" VARCHAR(100),
+        name VARCHAR(255) NOT NULL,
+        icon VARCHAR(50),
+        sort_order INT NOT NULL DEFAULT 0,
+        content_type VARCHAR(20) NOT NULL CHECK (content_type IN ('agent', 'html', 'custom')),
+        content_config JSONB NOT NULL DEFAULT '{}',
+        enabled BOOLEAN NOT NULL DEFAULT true,
+        created_at TIMESTAMP NOT NULL DEFAULT NOW(),
+        updated_at TIMESTAMP NOT NULL DEFAULT NOW()
+      )
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_menu_items_tenant_target
+      ON lattice_menu_items(tenant_id, menu_target)
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_menu_items_sort
+      ON lattice_menu_items(tenant_id, menu_target, sort_order)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query("DROP TABLE IF EXISTS lattice_menu_items");
+  },
+};

package/src/migrations/task_migration.ts

@@ -0,0 +1,51 @@
+import { PoolClient } from "pg";
+import { Migration } from "./migration";
+
+export const createTasksTable: Migration = {
+  version: 133,
+  name: "create_lattice_tasks_table",
+  up: async (client: PoolClient) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_tasks (
+        id VARCHAR(255) PRIMARY KEY,
+        tenant_id VARCHAR(255) NOT NULL,
+        owner_type VARCHAR(10) NOT NULL DEFAULT 'user',
+        owner_id VARCHAR(255) NOT NULL,
+        title VARCHAR(500) NOT NULL,
+        description TEXT,
+        status VARCHAR(20) NOT NULL DEFAULT 'pending',
+        priority VARCHAR(10) NOT NULL DEFAULT 'medium',
+        due_date VARCHAR(50),
+        metadata JSONB,
+        parent_id VARCHAR(255),
+        source_id VARCHAR(255),
+        context JSONB,
+        created_at VARCHAR(50) NOT NULL,
+        updated_at VARCHAR(50) NOT NULL
+      )
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_tasks_tenant_id
+      ON lattice_tasks(tenant_id)
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_tasks_owner
+      ON lattice_tasks(tenant_id, owner_type, owner_id)
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_tasks_status
+      ON lattice_tasks(tenant_id, status)
+    `);
+
+    await client.query(`
+      CREATE INDEX IF NOT EXISTS idx_lattice_tasks_parent_id
+      ON lattice_tasks(tenant_id, parent_id)
+    `);
+  },
+  down: async (client: PoolClient) => {
+    await client.query("DROP TABLE IF EXISTS lattice_tasks");
+  },
+};

package/src/migrations/workflow_tracking_migrations.ts

@@ -77,3 +77,22 @@ export const createWorkflowTrackingTables: Migration = {
     `);
   },
 };
+
+export const addStepThreadId: Migration = {
+  name: "add_step_thread_id",
+  version: 113,
+
+  up: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_workflow_steps
+      ADD COLUMN IF NOT EXISTS thread_id VARCHAR(255);
+    `);
+  },
+
+  down: async (client) => {
+    await client.query(`
+      ALTER TABLE lattice_workflow_steps
+      DROP COLUMN IF EXISTS thread_id;
+    `);
+  },
+};

package/src/stores/MenuStore.ts

@@ -0,0 +1,204 @@
+import { Pool } from "pg";
+import type { PoolConfig } from "pg";
+import type {
+  MenuItem,
+  MenuRegistry,
+  CreateMenuItemInput,
+  UpdateMenuItemInput,
+  MenuContentConfig,
+} from "@axiom-lattice/protocols";
+import { MigrationManager } from "../migrations/migration";
+import { createMenuItemsTable } from "../migrations/menu_items_migration";
+
+export interface MenuStoreOptions {
+  poolConfig: string | PoolConfig;
+  autoMigrate?: boolean;
+}
+
+type MenuItemRow = {
+  id: string;
+  tenant_id: string;
+  menu_target: string;
+  group: string | null;
+  name: string;
+  icon: string | null;
+  sort_order: number;
+  content_type: string;
+  content_config: Record<string, unknown>;
+  enabled: boolean;
+  created_at: Date;
+  updated_at: Date;
+};
+
+export class MenuStore implements MenuRegistry {
+  private pool: Pool;
+  private migrationManager: MigrationManager;
+  private initialized = false;
+  private initPromise: Promise<void> | null = null;
+
+  constructor(options: MenuStoreOptions) {
+    this.pool =
+      typeof options.poolConfig === "string"
+        ? new Pool({ connectionString: options.poolConfig })
+        : new Pool(options.poolConfig);
+
+    this.migrationManager = new MigrationManager(this.pool);
+    this.migrationManager.register(createMenuItemsTable);
+
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error("Failed to initialize MenuStore:", 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 list(params: { tenantId: string; menuTarget?: string }): Promise<MenuItem[]> {
+    await this.ensureInitialized();
+    const conditions: string[] = ["tenant_id = $1"];
+    const values: unknown[] = [params.tenantId];
+    let idx = 2;
+
+    if (params.menuTarget) {
+      conditions.push(`menu_target = $${idx++}`);
+      values.push(params.menuTarget);
+    }
+
+    const result = await this.pool.query<MenuItemRow>(
+      `SELECT * FROM lattice_menu_items
+       WHERE ${conditions.join(" AND ")}
+         AND enabled = true
+       ORDER BY sort_order ASC`,
+      values,
+    );
+    return result.rows.map((r) => this.mapRowToMenuItem(r));
+  }
+
+  async getById(id: string): Promise<MenuItem | null> {
+    await this.ensureInitialized();
+    const result = await this.pool.query<MenuItemRow>(
+      `SELECT * FROM lattice_menu_items WHERE id = $1`,
+      [id],
+    );
+    if (result.rows.length === 0) return null;
+    return this.mapRowToMenuItem(result.rows[0]);
+  }
+
+  async create(input: CreateMenuItemInput & { tenantId: string }): Promise<MenuItem> {
+    await this.ensureInitialized();
+    const result = await this.pool.query<MenuItemRow>(
+      `INSERT INTO lattice_menu_items
+         (tenant_id, menu_target, "group", name, icon, sort_order, content_type, content_config)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+       RETURNING *`,
+      [
+        input.tenantId,
+        input.menuTarget,
+        input.group || null,
+        input.name,
+        input.icon || null,
+        input.sortOrder ?? 0,
+        input.contentType,
+        JSON.stringify(input.contentConfig),
+      ],
+    );
+    return this.mapRowToMenuItem(result.rows[0]);
+  }
+
+  async update(id: string, patch: UpdateMenuItemInput): Promise<MenuItem> {
+    await this.ensureInitialized();
+    const existing = await this.pool.query<MenuItemRow>(
+      `SELECT * FROM lattice_menu_items WHERE id = $1`,
+      [id],
+    );
+    if (existing.rows.length === 0) {
+      throw new Error(`MenuItem ${id} not found`);
+    }
+
+    const row = existing.rows[0];
+    const setClauses: string[] = [];
+    const values: unknown[] = [];
+    let idx = 1;
+
+    if (patch.group !== undefined) {
+      setClauses.push(`"group" = $${idx++}`);
+      values.push(patch.group);
+    }
+    if (patch.name !== undefined) {
+      setClauses.push(`name = $${idx++}`);
+      values.push(patch.name);
+    }
+    if (patch.icon !== undefined) {
+      setClauses.push(`icon = $${idx++}`);
+      values.push(patch.icon);
+    }
+    if (patch.sortOrder !== undefined) {
+      setClauses.push(`sort_order = $${idx++}`);
+      values.push(patch.sortOrder);
+    }
+    if (patch.contentConfig !== undefined) {
+      setClauses.push(`content_config = $${idx++}`);
+      values.push(JSON.stringify(patch.contentConfig));
+    }
+    if (patch.enabled !== undefined) {
+      setClauses.push(`enabled = $${idx++}`);
+      values.push(patch.enabled);
+    }
+
+    if (setClauses.length === 0) {
+      return this.mapRowToMenuItem(row);
+    }
+
+    setClauses.push(`updated_at = NOW()`);
+    values.push(id);
+
+    const result = await this.pool.query<MenuItemRow>(
+      `UPDATE lattice_menu_items SET ${setClauses.join(", ")} WHERE id = $${idx}
+       RETURNING *`,
+      values,
+    );
+    return this.mapRowToMenuItem(result.rows[0]);
+  }
+
+  async delete(id: string): Promise<void> {
+    await this.ensureInitialized();
+    await this.pool.query(`DELETE FROM lattice_menu_items WHERE id = $1`, [id]);
+  }
+
+  private async ensureInitialized(): Promise<void> {
+    if (!this.initialized) {
+      await this.initialize();
+    }
+  }
+
+  private mapRowToMenuItem(row: MenuItemRow): MenuItem {
+    return {
+      id: row.id,
+      tenantId: row.tenant_id,
+      menuTarget: row.menu_target as MenuItem["menuTarget"],
+      group: row.group || undefined,
+      name: row.name,
+      icon: row.icon || undefined,
+      sortOrder: row.sort_order,
+      contentType: row.content_type as MenuItem["contentType"],
+      contentConfig: row.content_config as unknown as MenuContentConfig,
+      enabled: row.enabled,
+      createdAt: row.created_at,
+      updatedAt: row.updated_at,
+    };
+  }
+}

package/src/stores/PostgreSQLAssistantStore.ts

@@ -13,6 +13,7 @@ import { MigrationManager } from "../migrations/migration";
 import { createAssistantsTable } from "../migrations/assistant_migrations";
 import { addAssistantTenantId } from "../migrations/assistant_tenant_migration";
 import { changeAssistantPrimaryKey } from "../migrations/assistant_pk_migration";
+import { addAssistantOwnerUserId } from "../migrations/assistant_owner_user_migration";
 
 /**
  * PostgreSQL AssistantStore options
@@ -56,6 +57,7 @@ export class PostgreSQLAssistantStore implements AssistantStore {
     this.migrationManager.register(createAssistantsTable);
     this.migrationManager.register(addAssistantTenantId);
     this.migrationManager.register(changeAssistantPrimaryKey);
+    this.migrationManager.register(addAssistantOwnerUserId);
 
     // Auto-migrate by default
     if (options.autoMigrate !== false) {
@@ -103,11 +105,12 @@ export class PostgreSQLAssistantStore implements AssistantStore {
       name: string;
       description: string | null;
       graph_definition: any;
+      owner_user_id: string | null;
       created_at: Date;
       updated_at: Date;
     }>(
       `
-      SELECT id, tenant_id, name, description, graph_definition, created_at, updated_at
+      SELECT id, tenant_id, name, description, graph_definition, owner_user_id, created_at, updated_at
       FROM lattice_assistants
       WHERE tenant_id = $1
       ORDER BY created_at DESC
@@ -130,11 +133,12 @@ export class PostgreSQLAssistantStore implements AssistantStore {
       name: string;
       description: string | null;
       graph_definition: any;
+      owner_user_id: string | null;
       created_at: Date;
       updated_at: Date;
     }>(
       `
-      SELECT id, tenant_id, name, description, graph_definition, created_at, updated_at
+      SELECT id, tenant_id, name, description, graph_definition, owner_user_id, created_at, updated_at
       FROM lattice_assistants
       WHERE tenant_id = $1 AND id = $2
     `,
@@ -162,12 +166,13 @@ export class PostgreSQLAssistantStore implements AssistantStore {
 
     await this.pool.query(
       `
-      INSERT INTO lattice_assistants (id, tenant_id, name, description, graph_definition, created_at, updated_at)
-      VALUES ($1, $2, $3, $4, $5, $6, $7)
+      INSERT INTO lattice_assistants (id, tenant_id, name, description, graph_definition, owner_user_id, created_at, updated_at)
+      VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
       ON CONFLICT (id, tenant_id) DO UPDATE SET
         name = EXCLUDED.name,
         description = EXCLUDED.description,
         graph_definition = EXCLUDED.graph_definition,
+        owner_user_id = EXCLUDED.owner_user_id,
         updated_at = EXCLUDED.updated_at
     `,
       [
@@ -176,6 +181,7 @@ export class PostgreSQLAssistantStore implements AssistantStore {
         data.name,
         data.description || null,
         JSON.stringify(data.graphDefinition),
+        data.ownerUserId || null,
         now,
         now,
       ]
@@ -187,6 +193,7 @@ export class PostgreSQLAssistantStore implements AssistantStore {
       name: data.name,
       description: data.description,
       graphDefinition: data.graphDefinition,
+      ownerUserId: data.ownerUserId,
       createdAt: now,
       updatedAt: now,
     };
@@ -228,6 +235,11 @@ export class PostgreSQLAssistantStore implements AssistantStore {
       updateValues.push(JSON.stringify(updates.graphDefinition));
     }
 
+    if (updates.ownerUserId !== undefined) {
+      updateFields.push(`owner_user_id = $${paramIndex++}`);
+      updateValues.push(updates.ownerUserId || null);
+    }
+
     if (updateFields.length === 0) {
       // No fields to update
       return existing;
@@ -289,6 +301,35 @@ export class PostgreSQLAssistantStore implements AssistantStore {
     return result.rows.length > 0;
   }
 
+  /**
+   * Get assistant by owner user ID
+   */
+  async getByOwner(tenantId: string, userId: string): Promise<Assistant | null> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<{
+      id: string;
+      tenant_id: string;
+      name: string;
+      description: string | null;
+      graph_definition: any;
+      owner_user_id: string | null;
+      created_at: Date;
+      updated_at: Date;
+    }>(
+      `
+      SELECT id, tenant_id, name, description, graph_definition, owner_user_id, created_at, updated_at
+      FROM lattice_assistants
+      WHERE tenant_id = $1 AND owner_user_id = $2
+      LIMIT 1
+    `,
+      [tenantId, userId]
+    );
+
+    if (result.rows.length === 0) return null;
+    return this.mapRowToAssistant(result.rows[0]);
+  }
+
   /**
    * Dispose resources and close the connection pool
    * Should be called when the store is no longer needed
@@ -317,6 +358,7 @@ export class PostgreSQLAssistantStore implements AssistantStore {
     name: string;
     description: string | null;
     graph_definition: any;
+    owner_user_id: string | null;
     created_at: Date;
     updated_at: Date;
   }): Assistant {
@@ -329,6 +371,7 @@ export class PostgreSQLAssistantStore implements AssistantStore {
         typeof row.graph_definition === "string"
           ? JSON.parse(row.graph_definition)
           : row.graph_definition,
+      ownerUserId: row.owner_user_id || undefined,
       createdAt: row.created_at,
       updatedAt: row.updated_at,
     };