@kernl-sdk/pg 0.1.9

package/src/__tests__/thread.test.ts

@@ -0,0 +1,285 @@
+import { describe, it, expect, beforeAll, afterAll, beforeEach } from "vitest";
+import { Pool } from "pg";
+
+import {
+  Agent,
+  Kernl,
+  tool,
+  FunctionToolkit,
+  KernlStorage,
+} from "kernl";
+import { STOPPED, message, IN_PROGRESS } from "@kernl-sdk/protocol";
+
+import { postgres } from "@/postgres";
+
+// Use helper from test fixtures (handles streaming properly)
+function createMockModel(generateFn: any): any {
+  return {
+    spec: "1.0" as const,
+    provider: "test",
+    modelId: "test-model",
+    generate: generateFn,
+    stream: async function* (req: any) {
+      const response = await generateFn(req);
+      for (const item of response.content) {
+        yield item;
+      }
+      yield {
+        kind: "finish" as const,
+        finishReason: response.finishReason,
+        usage: response.usage,
+      };
+    },
+  };
+}
+
+const TEST_DB_URL = process.env.KERNL_PG_TEST_URL;
+
+describe.sequential("PG Thread Lifecycle", () => {
+  if (!TEST_DB_URL) {
+    it.skip("requires KERNL_PG_TEST_URL environment variable", () => {});
+    return;
+  }
+
+  let pool: Pool;
+  let storage: KernlStorage;
+  let kernl: Kernl;
+
+  beforeAll(async () => {
+    pool = new Pool({ connectionString: TEST_DB_URL });
+    storage = postgres({ pool });
+
+    await pool.query('DROP SCHEMA IF EXISTS "kernl" CASCADE');
+    await storage.init();
+
+    kernl = new Kernl({ storage: { db: storage } });
+  });
+
+  afterAll(async () => {
+    await storage.close();
+  });
+
+  beforeEach(async () => {
+    // Clean threads between tests
+    await pool.query(
+      'TRUNCATE "kernl"."threads", "kernl"."thread_events" CASCADE',
+    );
+  });
+
+  describe("Simple thread (no tools)", () => {
+    it("should persist thread record and events correctly", async () => {
+      const model = createMockModel(async () => ({
+        content: [message({ role: "assistant", text: "Hello, world!" })],
+        finishReason: "stop",
+        usage: { inputTokens: 5, outputTokens: 3, totalTokens: 8 },
+        warnings: [],
+      }));
+
+      const agent = new Agent({
+        id: "simple-agent",
+        name: "Simple",
+        instructions: "Test",
+        model,
+      });
+
+      kernl.register(agent);
+
+      const result = await agent.run("Hello");
+
+      expect(result.response).toBe("Hello, world!");
+
+      // Verify thread record
+      const threadResult = await pool.query(
+        'SELECT * FROM "kernl"."threads" WHERE agent_id = $1',
+        ["simple-agent"],
+      );
+
+      expect(threadResult.rows).toHaveLength(1);
+      const thread = threadResult.rows[0];
+
+      expect(thread.agent_id).toBe("simple-agent");
+      expect(thread.state).toBe(STOPPED);
+      expect(thread.tick).toBe(1);
+      expect(thread.model).toMatch(/test/);
+
+      // Verify events
+      const eventsResult = await pool.query(
+        'SELECT * FROM "kernl"."thread_events" WHERE tid = $1 ORDER BY seq ASC',
+        [thread.id],
+      );
+
+      expect(eventsResult.rows.length).toBeGreaterThanOrEqual(2);
+
+      // User message
+      expect(eventsResult.rows[0].kind).toBe("message");
+      expect(eventsResult.rows[0].seq).toBe(0);
+
+      // Assistant message
+      expect(eventsResult.rows[1].kind).toBe("message");
+      expect(eventsResult.rows[1].seq).toBe(1);
+
+      // Verify monotonic seq
+      const seqs = eventsResult.rows.map((r) => r.seq);
+      expect(seqs).toEqual([...seqs].sort((a, b) => a - b));
+      expect(new Set(seqs).size).toBe(seqs.length);
+    });
+  });
+
+  describe("Multi-turn with tools", () => {
+    it("should persist all events across ticks with correct seq", async () => {
+      let callCount = 0;
+      const model = createMockModel(async () => {
+        callCount++;
+        if (callCount === 1) {
+          return {
+            content: [
+              message({ role: "assistant", text: "" }),
+              {
+                kind: "tool-call",
+                toolId: "add",
+                callId: "call_1",
+                state: IN_PROGRESS,
+                arguments: '{"a":2,"b":3}',
+              },
+            ],
+            finishReason: "stop",
+            usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 },
+            warnings: [],
+          };
+        }
+        return {
+          content: [message({ role: "assistant", text: "The sum is 5" })],
+          finishReason: "stop",
+          usage: { inputTokens: 15, outputTokens: 5, totalTokens: 20 },
+          warnings: [],
+        };
+      });
+
+      const addTool = tool({
+        id: "add",
+        description: "Add numbers",
+        parameters: undefined, // No validation for test
+        execute: async (ctx, params: any) => {
+          const { a, b } = params;
+          return a + b;
+        },
+      });
+
+      const agent = new Agent({
+        id: "tool-agent",
+        name: "Tool Agent",
+        instructions: "Test",
+        model,
+        toolkits: [new FunctionToolkit({ id: "math", tools: [addTool] })],
+      });
+
+      kernl.register(agent);
+
+      await agent.run("Add 2 and 3");
+
+      // verify thread
+      const threadResult = await pool.query(
+        'SELECT * FROM "kernl"."threads" WHERE agent_id = $1',
+        ["tool-agent"],
+      );
+
+      expect(threadResult.rows).toHaveLength(1);
+      expect(threadResult.rows[0].tick).toBe(2); // 2 ticks
+      expect(threadResult.rows[0].state).toBe(STOPPED);
+
+      // verify events
+      const eventsResult = await pool.query(
+        'SELECT * FROM "kernl"."thread_events" WHERE tid = $1 ORDER BY seq ASC',
+        [threadResult.rows[0].id],
+      );
+
+      const events = eventsResult.rows;
+
+      // Should have: user msg, assistant msg (tick1), tool-call, tool-result, assistant msg (tick2)
+      expect(events.length).toBeGreaterThanOrEqual(5);
+
+      // Verify event kinds and order
+      expect(events[0].kind).toBe("message"); // user
+      expect(events[1].kind).toBe("message"); // assistant (tick 1)
+      expect(events[2].kind).toBe("tool-call");
+      expect(events[3].kind).toBe("tool-result");
+      expect(events[4].kind).toBe("message"); // assistant (tick 2)
+
+      // Verify tool result exists and is in correct format
+      const toolResult = events[3];
+      expect(toolResult.data.state).toBe("completed");
+      expect(toolResult.data.callId).toBe("call_1");
+      expect(toolResult.data.toolId).toBe("add");
+
+      // Verify seq is monotonic
+      const seqs = events.map((e) => e.seq);
+      expect(seqs).toEqual([0, 1, 2, 3, 4]);
+    });
+  });
+
+  describe("Resume from storage", () => {
+    it("should load thread from storage and append new events", async () => {
+      const model = createMockModel(async () => ({
+        content: [message({ role: "assistant", text: "Response" })],
+        finishReason: "stop",
+        usage: { inputTokens: 5, outputTokens: 3, totalTokens: 8 },
+        warnings: [],
+      }));
+
+      const agent = new Agent({
+        id: "resume-agent",
+        name: "Resume",
+        instructions: "Test",
+        model,
+      });
+
+      kernl.register(agent);
+
+      const threadId = "resume-test-123";
+
+      // First run
+      await agent.run("First message", { threadId });
+
+      const firstEventsResult = await pool.query(
+        'SELECT COUNT(*) as count FROM "kernl"."thread_events" WHERE tid = $1',
+        [threadId],
+      );
+
+      const firstEventCount = parseInt(firstEventsResult.rows[0].count);
+      expect(firstEventCount).toBeGreaterThanOrEqual(2);
+
+      // Second run: resume
+      await agent.run("Second message", { threadId });
+
+      const secondEventsResult = await pool.query(
+        'SELECT * FROM "kernl"."thread_events" WHERE tid = $1 ORDER BY seq ASC',
+        [threadId],
+      );
+
+      const secondEventCount = secondEventsResult.rows.length;
+
+      // Should have more events now
+      expect(secondEventCount).toBeGreaterThan(firstEventCount);
+
+      // Verify seq is still monotonic across both runs
+      const seqs = secondEventsResult.rows.map((r) => r.seq);
+      expect(seqs).toEqual([...seqs].sort((a, b) => a - b));
+      expect(new Set(seqs).size).toBe(seqs.length);
+
+      // Verify last events are from second run
+      const lastEvents = secondEventsResult.rows.slice(-2);
+      expect(lastEvents[0].kind).toBe("message"); // user: "Second message"
+      expect(lastEvents[1].kind).toBe("message"); // assistant response
+    });
+  });
+
+  describe("Error handling", () => {
+    it.skip("should rollback on persist failure", async () => {
+      // TODO: Implement once we have transaction-level control in storage
+      // This would require forcing a DB error mid-transaction
+      // Expected behavior:
+      // - If append() fails, no partial events should be written
+      // - Thread state should remain consistent
+    });
+  });
+});

package/src/index.ts

@@ -0,0 +1,7 @@
+/**
+ * @kernl/pg - PostgreSQL storage adapter for Kernl
+ */
+
+export { PGStorage, type PGStorageConfig } from "./storage";
+export { postgres, type PostgresConfig } from "./postgres";
+export { migrations, REQUIRED_SCHEMA_VERSION } from "./migrations";

package/src/migrations.ts

@@ -0,0 +1,38 @@
+/**
+ * Database migrations.
+ */
+
+import type { PoolClient } from "pg";
+import type { Table, Column } from "@kernl-sdk/storage";
+import { TABLE_THREADS, TABLE_THREAD_EVENTS, SCHEMA_NAME } from "@kernl-sdk/storage";
+
+/**
+ * Migration context with helpers.
+ */
+export interface MigrationContext {
+  client: PoolClient;
+  createTable: (table: Table<string, Record<string, Column>>) => Promise<void>;
+}
+
+export interface Migration {
+  id: string;
+  up: (ctx: MigrationContext) => Promise<void>;
+}
+
+/**
+ * List of all migrations in order.
+ */
+export const migrations: Migration[] = [
+  {
+    id: "0001_initial",
+    async up(ctx) {
+      await ctx.createTable(TABLE_THREADS);
+      await ctx.createTable(TABLE_THREAD_EVENTS);
+    },
+  },
+];
+
+/**
+ * Minimum schema version required by this version of @kernl/pg.
+ */
+export const REQUIRED_SCHEMA_VERSION = "0001_initial";

package/src/postgres.ts

@@ -0,0 +1,63 @@
+import { Pool } from "pg";
+import type { KernlStorage } from "kernl";
+
+import { PGStorage } from "./storage";
+
+/**
+ * PostgreSQL connection configuration.
+ */
+export type PostgresConfig =
+  | { pool: Pool }
+  | { connstr: string }
+  | {
+      host: string;
+      port: number;
+      database: string;
+      user: string;
+      password: string;
+    };
+
+/**
+ * Create a PostgreSQL storage adapter for Kernl.
+ *
+ * @param config - Connection configuration (pool, connection string, or credentials)
+ * @returns KernlStorage instance backed by PostgreSQL
+ *
+ * @example
+ * ```ts
+ * // with connection string
+ * const storage = postgres({ connstr: "postgresql://localhost/mydb" });
+ *
+ * // with connection options
+ * const storage = postgres({
+ *   host: "localhost",
+ *   port: 5432,
+ *   database: "mydb",
+ *   user: "user",
+ *   password: "password"
+ * });
+ *
+ * // existing pool
+ * const pool = new Pool({ ... });
+ * const storage = postgres({ pool });
+ * ```
+ */
+export function postgres(config: PostgresConfig): KernlStorage {
+  let pool: Pool;
+
+  if ("pool" in config) {
+    pool = config.pool;
+  } else if ("connstr" in config) {
+    pool = new Pool({ connectionString: config.connstr });
+  } else {
+    pool = new Pool({
+      host: config.host,
+      port: config.port,
+      database: config.database,
+      user: config.user,
+      password: config.password,
+    });
+  }
+
+  return new PGStorage({ pool });
+}

package/src/sql.ts

@@ -0,0 +1,8 @@
+/**
+ * SQL utilities for safe query construction.
+ */
+
+/**
+ * SQL identifier regex - alphanumeric + underscore, must start with letter/underscore.
+ */
+export const SQL_IDENTIFIER_REGEX = /^[a-zA-Z_][a-zA-Z0-9_]*$/;

package/src/storage.ts

@@ -0,0 +1,270 @@
+import assert from "assert";
+import type { Pool, PoolClient } from "pg";
+
+/* workspace */
+import type { Table, Column, IndexConstraint } from "@kernl-sdk/storage";
+import { SCHEMA_NAME, TABLE_MIGRATIONS } from "@kernl-sdk/storage";
+import type {
+  AgentRegistry,
+  ModelRegistry,
+  KernlStorage,
+  Transaction,
+} from "kernl";
+import { UnimplementedError } from "@kernl-sdk/shared/lib";
+
+/* pg */
+import { PGThreadStore } from "./thread/store";
+import { SQL_IDENTIFIER_REGEX } from "./sql";
+import { migrations } from "./migrations";
+
+/**
+ * PostgreSQL storage configuration.
+ */
+export interface PGStorageConfig {
+  /**
+   * Pool instance for database connections.
+   */
+  pool: Pool;
+}
+
+/**
+ * PostgreSQL storage adapter.
+ */
+export class PGStorage implements KernlStorage {
+  private pool: Pool;
+
+  threads: PGThreadStore;
+
+  constructor(config: PGStorageConfig) {
+    this.pool = config.pool;
+    this.threads = new PGThreadStore(this.pool);
+  }
+
+  /**
+   * Bind runtime registries to storage.
+   */
+  bind(registries: { agents: AgentRegistry; models: ModelRegistry }): void {
+    this.threads.bind(registries);
+  }
+
+  /**
+   * Execute a function within a transaction.
+   */
+  async transaction<T>(fn: (tx: Transaction) => Promise<T>): Promise<T> {
+    throw new UnimplementedError();
+  }
+
+  /**
+   * Initialize the storage backend.
+   */
+  async init(): Promise<void> {
+    await this.pool.query(`CREATE SCHEMA IF NOT EXISTS "${SCHEMA_NAME}"`);
+    await this.createTable(TABLE_MIGRATIONS);
+    await this.migrate();
+  }
+
+  /**
+   * Close the storage backend and cleanup resources.
+   */
+  async close(): Promise<void> {
+    await this.pool.end();
+  }
+
+  /**
+   * Run migrations to ensure all required tables exist.
+   */
+  async migrate(): Promise<void> {
+    const client = await this.pool.connect();
+    try {
+      await client.query("BEGIN");
+
+      // read applied migration IDs
+      const result = await client.query<{ id: string }>(
+        `SELECT id FROM "${SCHEMA_NAME}".migrations ORDER BY applied_at ASC`,
+      );
+      const applied = new Set(result.rows.map((row) => row.id));
+
+      // filter pending migrations
+      const pending = migrations.filter((m) => !applied.has(m.id));
+      if (pending.length === 0) {
+        await client.query("COMMIT");
+        return;
+      }
+
+      // run pending migrations + insert into migrations table
+      for (const migration of pending) {
+        await migration.up({
+          client,
+          createTable: async (table: Table<string, Record<string, Column>>) => {
+            await this._createTable(client, table);
+          },
+        });
+        await client.query(
+          `INSERT INTO "${SCHEMA_NAME}".migrations (id, applied_at) VALUES ($1, $2)`,
+          [migration.id, Date.now()],
+        );
+      }
+
+      await client.query("COMMIT");
+    } catch (error) {
+      await client.query("ROLLBACK");
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
+  /**
+   * Create a table from its definition.
+   */
+  private async createTable(
+    table: Table<string, Record<string, Column>>,
+  ): Promise<void> {
+    assert(
+      SQL_IDENTIFIER_REGEX.test(table.name),
+      "system table should have a valid name",
+    );
+    await this._createTable(this.pool, table);
+  }
+
+  /**
+   * Create a table from its definition using a specific client.
+   */
+  private async _createTable(
+    client: Pool | PoolClient,
+    table: Table<string, Record<string, Column>>,
+  ): Promise<void> {
+    const columns: string[] = [];
+    const tableConstraints: string[] = [];
+
+    // build column definitions
+    for (const name in table.columns) {
+      const col = table.columns[name];
+
+      const constraints: string[] = [];
+      if (col._pk) constraints.push("PRIMARY KEY");
+      if (col._unique) constraints.push("UNIQUE");
+      if (!col._nullable && !col._pk) constraints.push("NOT NULL");
+      if (col._default !== undefined) {
+        constraints.push(`DEFAULT ${col.encode(col._default)}`);
+      }
+
+      // foreign key reference
+      if (col._fk) {
+        let ref = `REFERENCES "${SCHEMA_NAME}"."${col._fk.table}" ("${col._fk.column}")`;
+        if (col._onDelete) {
+          ref += ` ON DELETE ${col._onDelete}`;
+        }
+        constraints.push(ref);
+      }
+
+      columns.push(
+        `"${name}" ${col.type.toUpperCase()} ${constraints.join(" ")}`.trim(),
+      );
+    }
+
+    // table-level constraints
+    const indexes: IndexConstraint[] = [];
+    if (table.constraints) {
+      for (const constraint of table.constraints) {
+        switch (constraint.kind) {
+          case "unique": {
+            const name =
+              constraint.name ??
+              `${table.name}_${constraint.columns.join("_")}_unique`;
+            const cols = constraint.columns.map((c) => `"${c}"`).join(", ");
+            tableConstraints.push(`CONSTRAINT "${name}" UNIQUE (${cols})`);
+            break;
+          }
+
+          case "pkey": {
+            const name = constraint.name ?? `${table.name}_pkey`;
+            const cols = constraint.columns.map((c) => `"${c}"`).join(", ");
+            tableConstraints.push(`CONSTRAINT "${name}" PRIMARY KEY (${cols})`);
+            break;
+          }
+
+          case "fkey": {
+            throw new UnimplementedError(
+              "Composite foreign keys not yet supported. Use column-level .references() for single-column FKs.",
+            );
+          }
+
+          case "check": {
+            const name = constraint.name ?? `${table.name}_check`;
+            tableConstraints.push(
+              `CONSTRAINT "${name}" CHECK (${constraint.expression})`,
+            );
+            break;
+          }
+
+          case "index": {
+            // collect indexes to create after table
+            indexes.push(constraint);
+            break;
+          }
+        }
+      }
+    }
+
+    const constraints = [...columns, ...tableConstraints];
+
+    const sql = `
+      CREATE TABLE IF NOT EXISTS "${SCHEMA_NAME}"."${table.name}" (
+        ${constraints.join(",\n  ")}
+      )
+    `.trim();
+
+    await client.query(sql);
+
+    // create indexes
+    for (const index of indexes) {
+      await this.createIndex(client, table.name, index);
+    }
+  }
+
+  /**
+   * Alter a table definition (not implemented).
+   */
+  private async alterTable(
+    table: Table<string, Record<string, Column>>,
+  ): Promise<void> {
+    throw new UnimplementedError();
+  }
+
+  /**
+   * Clear all rows from a table (not implemented).
+   */
+  private async clearTable(tableName: string): Promise<void> {
+    throw new UnimplementedError();
+  }
+
+  /**
+   * Drop a table (not implemented).
+   */
+  private async dropTable(tableName: string): Promise<void> {
+    throw new UnimplementedError();
+  }
+
+  /**
+   * Create an index from its definition.
+   */
+  private async createIndex(
+    client: Pool | PoolClient,
+    tableName: string,
+    index: IndexConstraint,
+  ): Promise<void> {
+    const uniqueKeyword = index.unique ? "UNIQUE" : "";
+    const columns = index.columns.map((c) => `"${c}"`).join(", ");
+
+    // Auto-generate index name: idx_{table}_{col1}_{col2}...
+    const indexName = `idx_${tableName}_${index.columns.join("_")}`;
+
+    const sql = `
+      CREATE ${uniqueKeyword} INDEX IF NOT EXISTS "${indexName}"
+      ON "${SCHEMA_NAME}"."${tableName}" (${columns})
+    `.trim();
+
+    await client.query(sql);
+  }
+}