@ixo/sqlite-saver 1.0.4

package/src/index.ts

@@ -0,0 +1,929 @@
+import type { RunnableConfig } from '@langchain/core/runnables';
+import {
+  BaseCheckpointSaver,
+  type Checkpoint,
+  type CheckpointListOptions,
+  type CheckpointMetadata,
+  type CheckpointTuple,
+  copyCheckpoint,
+  maxChannelVersion,
+  type PendingWrite,
+  type SerializerProtocol,
+  TASKS,
+} from '@langchain/langgraph-checkpoint';
+import Database, { Database as DatabaseType, Statement } from 'better-sqlite3';
+import { BaseMessage } from 'langchain';
+import migration001 from './migrations/001_add_created_at_to_messages';
+import {
+  _default,
+  CleanAdditionalKwargs,
+  cleanAdditionalKwargs,
+  stringify,
+} from './utils';
+
+type ChannelValues<C extends string = string> = {
+  [key in C]: unknown;
+};
+
+interface CheckpointWithMessages<
+  N extends string = string,
+  C extends string = string,
+> extends Checkpoint<N, C> {
+  channel_values: ChannelValues<C> & {
+    messages?: BaseMessage[];
+  };
+}
+
+interface MessageRow {
+  thread_id: string;
+  checkpoint_ns: string;
+  checkpoint_id: string;
+  message_id: string;
+  message_type: string;
+  message_content: string;
+  message: string;
+  created_at: string;
+}
+
+interface CheckpointRow {
+  checkpoint: string;
+  metadata: string;
+  parent_checkpoint_id?: string;
+  thread_id: string;
+  checkpoint_id: string;
+  checkpoint_ns?: string;
+  type?: string;
+  pending_writes: string;
+}
+
+interface PendingWriteColumn {
+  task_id: string;
+  channel: string;
+  type: string;
+  value: string;
+}
+
+interface PendingSendColumn {
+  type: string;
+  value: string;
+}
+
+interface Migration {
+  version: number;
+  name: string;
+  up: (db: DatabaseType) => void;
+}
+
+// In the `SqliteSaver.list` method, we need to sanitize the `options.filter` argument to ensure it only contains keys
+// that are part of the `CheckpointMetadata` type. The lines below ensure that we get compile-time errors if the list
+// of keys that we use is out of sync with the `CheckpointMetadata` type.
+const checkpointMetadataKeys = ['source', 'step', 'parents'] as const;
+
+type CheckKeys<T, K extends readonly (keyof T)[]> = [K[number]] extends [
+  keyof T,
+]
+  ? [keyof T] extends [K[number]]
+    ? K
+    : never
+  : never;
+
+function validateKeys<T, K extends readonly (keyof T)[]>(
+  keys: CheckKeys<T, K>,
+): K {
+  return keys;
+}
+
+// If this line fails to compile, the list of keys that we use in the `SqliteSaver.list` method is out of sync with the
+// `CheckpointMetadata` type. In that case, just update `checkpointMetadataKeys` to contain all the keys in
+// `CheckpointMetadata`
+const validCheckpointMetadataKeys = validateKeys<
+  CheckpointMetadata,
+  typeof checkpointMetadataKeys
+>(checkpointMetadataKeys);
+
+function prepareSql(db: DatabaseType, checkpointId: boolean) {
+  const sql = `
+  SELECT
+    thread_id,
+    checkpoint_ns,
+    checkpoint_id,
+    parent_checkpoint_id,
+    type,
+    checkpoint,
+    metadata,
+    (
+      SELECT
+        json_group_array(
+          json_object(
+            'task_id', pw.task_id,
+            'channel', pw.channel,
+            'type', pw.type,
+            'value', CAST(pw.value AS TEXT)
+          )
+        )
+      FROM writes as pw
+      WHERE pw.thread_id = checkpoints.thread_id
+        AND pw.checkpoint_ns = checkpoints.checkpoint_ns
+        AND pw.checkpoint_id = checkpoints.checkpoint_id
+    ) as pending_writes,
+    (
+      SELECT
+        json_group_array(
+          json_object(
+            'type', ps.type,
+            'value', CAST(ps.value AS TEXT)
+          )
+        )
+      FROM writes as ps
+      WHERE ps.thread_id = checkpoints.thread_id
+        AND ps.checkpoint_ns = checkpoints.checkpoint_ns
+        AND ps.checkpoint_id = checkpoints.parent_checkpoint_id
+        AND ps.channel = '${TASKS}'
+      ORDER BY ps.idx
+    ) as pending_sends
+  FROM checkpoints
+  WHERE thread_id = ? AND checkpoint_ns = ? ${
+    checkpointId
+      ? 'AND checkpoint_id = ?'
+      : 'ORDER BY checkpoint_id DESC LIMIT 1'
+  }`;
+
+  return db.prepare(sql);
+}
+
+export class SqliteSaver extends BaseCheckpointSaver {
+  db: DatabaseType;
+
+  protected isSetup: boolean;
+
+  protected withoutCheckpoint: Statement;
+
+  protected withCheckpoint: Statement;
+
+  protected putCheckpointStmt: Statement;
+
+  protected putWritesStmt: Statement;
+
+  protected deleteCheckpointsStmt: Statement;
+  protected putMessageStmt: Statement;
+
+  protected getMessageStmt: Statement;
+
+  protected deleteWritesStmt: Statement;
+
+  constructor(db: DatabaseType, serde?: SerializerProtocol) {
+    super(serde);
+    this.db = db;
+    this.isSetup = false;
+  }
+
+  static fromConnString(connStringOrLocalPath: string): SqliteSaver {
+    return new SqliteSaver(new Database(connStringOrLocalPath));
+  }
+
+  static fromDatabase(db: DatabaseType, serde?: SerializerProtocol): SqliteSaver {
+    return new SqliteSaver(db, serde);
+  }
+
+  close(): void {
+    if (this.db.open) {
+      this.db.close();
+    }
+  }
+
+  /**
+   * Create schema_migrations table to track applied migrations
+   */
+  protected createSchemaMigrationsTable(): void {
+    this.db.exec(`
+      CREATE TABLE IF NOT EXISTS schema_migrations (
+        version INTEGER NOT NULL PRIMARY KEY,
+        name TEXT NOT NULL,
+        applied_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
+      );
+    `);
+  }
+
+  /**
+   * Get list of applied migration versions
+   */
+  protected getAppliedMigrations(): number[] {
+    try {
+      const rows = this.db
+        .prepare('SELECT version FROM schema_migrations ORDER BY version')
+        .all() as Array<{ version: number }>;
+      return rows.map((row) => row.version);
+    } catch {
+      // Table doesn't exist yet, return empty array
+      return [];
+    }
+  }
+
+  /**
+   * Record that a migration has been applied
+   */
+  protected recordMigration(migration: Migration): void {
+    this.db
+      .prepare(
+        'INSERT INTO schema_migrations (version, name, applied_at) VALUES (?, ?, CURRENT_TIMESTAMP)',
+      )
+      .run(migration.version, migration.name);
+  }
+
+  /**
+   * Load all migration files from the migrations folder
+   *
+   * Note: When adding new migrations:
+   * 1. Import them at the top of this file
+   * 2. Add them to the migrations array below
+   */
+  protected loadMigrations(): Migration[] {
+    const migrations: Migration[] = [
+      migration001,
+      // Add future migrations here:
+      // migration002,
+      // migration003,
+    ];
+
+    // Sort by version to ensure correct order
+    migrations.sort((a, b) => a.version - b.version);
+
+    return migrations;
+  }
+
+  /**
+   * Run pending migrations
+   */
+  protected runMigrations(): void {
+    this.createSchemaMigrationsTable();
+
+    const appliedVersions = this.getAppliedMigrations();
+    const allMigrations = this.loadMigrations();
+
+    const pendingMigrations = allMigrations.filter(
+      (migration) => !appliedVersions.includes(migration.version),
+    );
+
+    if (pendingMigrations.length === 0) {
+      return;
+    }
+
+    console.log(`Running ${pendingMigrations.length} pending migration(s)...`);
+
+    for (const migration of pendingMigrations) {
+      try {
+        console.log(
+          `Applying migration ${migration.version}: ${migration.name}`,
+        );
+        migration.up(this.db);
+        this.recordMigration(migration);
+        console.log(
+          `✓ Migration ${migration.version}: ${migration.name} applied successfully`,
+        );
+      } catch (error) {
+        console.error(
+          `✗ Failed to apply migration ${migration.version}: ${migration.name}`,
+          error,
+        );
+        throw error;
+      }
+    }
+  }
+
+  protected setup(): void {
+    if (this.isSetup) {
+      return;
+    }
+
+    // Enable WAL mode for concurrent read/write support and set busy timeout
+    this.db.pragma('journal_mode = WAL');
+    this.db.pragma('busy_timeout = 5000');
+
+    // Create base tables
+    this.db.exec(`
+CREATE TABLE IF NOT EXISTS checkpoints (
+  thread_id TEXT NOT NULL,
+  checkpoint_ns TEXT NOT NULL DEFAULT '',
+  checkpoint_id TEXT NOT NULL,
+  parent_checkpoint_id TEXT,
+  type TEXT,
+  checkpoint BLOB,
+  metadata BLOB,
+  PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
+);`);
+    this.db.exec(`
+CREATE TABLE IF NOT EXISTS writes (
+  thread_id TEXT NOT NULL,
+  checkpoint_ns TEXT NOT NULL DEFAULT '',
+  checkpoint_id TEXT NOT NULL,
+  task_id TEXT NOT NULL,
+  idx INTEGER NOT NULL,
+  channel TEXT NOT NULL,
+  type TEXT,
+  value BLOB,
+  PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
+);`);
+
+    // Create messages table
+    // For new databases, created_at is included
+    // For existing databases, the migration will add it if missing
+    this.db.exec(`
+CREATE TABLE IF NOT EXISTS messages (
+  thread_id TEXT NOT NULL,
+  checkpoint_ns TEXT NOT NULL DEFAULT '',
+  checkpoint_id TEXT NOT NULL,
+  message_id TEXT NOT NULL,
+  message_type TEXT NOT NULL,
+  message_content TEXT NOT NULL,
+  message BLOB,
+  created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+  PRIMARY KEY  (message_id)
+    );
+`);
+
+    // Create schema_migrations table and run migrations
+    this.runMigrations();
+
+    // Create indexes (after migrations have run)
+    this.db.exec(`
+CREATE INDEX IF NOT EXISTS idx_messages_thread_id 
+ON messages(thread_id);
+`);
+    this.db.exec(`
+CREATE INDEX IF NOT EXISTS idx_messages_checkpoint_id 
+ON messages(checkpoint_id);
+`);
+
+    this.db.exec(`
+CREATE INDEX IF NOT EXISTS idx_messages_lookup 
+ON messages(thread_id, checkpoint_ns, checkpoint_id);
+`);
+
+    // Create index on created_at (migration ensures column exists)
+    this.db.exec(`
+CREATE INDEX IF NOT EXISTS idx_messages_thread_created 
+ON messages(thread_id, created_at);
+`);
+
+    this.db.exec(`
+CREATE INDEX IF NOT EXISTS idx_writes_channel 
+ON writes(thread_id, checkpoint_id, channel);
+`);
+
+    this.withoutCheckpoint = prepareSql(this.db, false);
+    this.withCheckpoint = prepareSql(this.db, true);
+
+    // Cache prepared statements for write operations
+    this.putCheckpointStmt = this.db.prepare(
+      `INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)`,
+    );
+    this.putWritesStmt = this.db.prepare(`
+      INSERT OR REPLACE INTO writes 
+      (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) 
+      VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+    `);
+    this.deleteCheckpointsStmt = this.db.prepare(
+      `DELETE FROM checkpoints WHERE thread_id = ?`,
+    );
+    this.deleteWritesStmt = this.db.prepare(
+      `DELETE FROM writes WHERE thread_id = ?`,
+    );
+
+    this.putMessageStmt = this.db.prepare(
+      `INSERT OR REPLACE INTO messages (thread_id, checkpoint_ns, checkpoint_id, message_id, message_type, message_content, message, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
+    );
+
+    this.getMessageStmt = this.db.prepare(
+      `SELECT * FROM messages WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?`,
+    );
+
+    this.isSetup = true;
+  }
+
+  async getTuple(config: RunnableConfig): Promise<CheckpointTuple | undefined> {
+    this.setup();
+    const {
+      thread_id,
+      checkpoint_ns = '',
+      checkpoint_id,
+    } = config.configurable ?? {};
+
+    const args = [thread_id, checkpoint_ns];
+    if (checkpoint_id) args.push(checkpoint_id);
+
+    const stm = checkpoint_id ? this.withCheckpoint : this.withoutCheckpoint;
+    const row = stm.get(...args) as CheckpointRow;
+    if (row === undefined) return undefined;
+
+    let finalConfig = config;
+
+    if (!checkpoint_id) {
+      finalConfig = {
+        configurable: {
+          thread_id: row.thread_id,
+          checkpoint_ns,
+          checkpoint_id: row.checkpoint_id,
+        },
+      };
+    }
+
+    if (
+      finalConfig.configurable?.thread_id === undefined ||
+      finalConfig.configurable?.checkpoint_id === undefined
+    ) {
+      throw new Error('Missing thread_id or checkpoint_id');
+    }
+
+    const messages = this.getMessageStmt.all(
+      finalConfig.configurable?.thread_id,
+      finalConfig.configurable?.checkpoint_ns,
+      finalConfig.configurable?.checkpoint_id,
+    ) as MessageRow[];
+
+    const pendingWrites = await Promise.all(
+      (JSON.parse(row.pending_writes) as PendingWriteColumn[]).map(
+        async (write) => {
+          return [
+            write.task_id,
+            write.channel,
+            await this.serde.loadsTyped(
+              write.type ?? 'json',
+              write.value ?? '',
+            ),
+          ] as [string, string, unknown];
+        },
+      ),
+    );
+
+    const parsedMessages: BaseMessage[] = await Promise.all(
+      messages.map(async (message) => {
+        return this.serde.loadsTyped('json', message.message);
+      }),
+    );
+
+    const checkpoint = (await this.serde.loadsTyped(
+      row.type ?? 'json',
+      row.checkpoint,
+    )) as Checkpoint;
+
+    if (parsedMessages.length > 0) {
+      checkpoint.channel_values.messages = parsedMessages;
+    }
+
+    if (checkpoint.v < 4 && row.parent_checkpoint_id != null) {
+      await this.migratePendingSends(
+        checkpoint,
+        row.thread_id,
+        row.parent_checkpoint_id,
+      );
+    }
+
+    return {
+      checkpoint,
+      config: finalConfig,
+      metadata: (await this.serde.loadsTyped(
+        row.type ?? 'json',
+        row.metadata,
+      )) as CheckpointMetadata,
+      parentConfig: row.parent_checkpoint_id
+        ? {
+            configurable: {
+              thread_id: row.thread_id,
+              checkpoint_ns,
+              checkpoint_id: row.parent_checkpoint_id,
+            },
+          }
+        : undefined,
+      pendingWrites,
+    };
+  }
+
+  async *list(
+    config: RunnableConfig,
+    options?: CheckpointListOptions,
+  ): AsyncGenerator<CheckpointTuple> {
+    const { limit, before, filter } = options ?? {};
+    this.setup();
+    const thread_id = config.configurable?.thread_id;
+    const checkpoint_ns = config.configurable?.checkpoint_ns;
+    let sql = `
+      SELECT
+        thread_id,
+        checkpoint_ns,
+        checkpoint_id,
+        parent_checkpoint_id,
+        type,
+        checkpoint,
+        metadata,
+        (
+          SELECT
+            json_group_array(
+              json_object(
+                'task_id', pw.task_id,
+                'channel', pw.channel,
+                'type', pw.type,
+                'value', CAST(pw.value AS TEXT)
+              )
+            )
+          FROM writes as pw
+          WHERE pw.thread_id = checkpoints.thread_id
+            AND pw.checkpoint_ns = checkpoints.checkpoint_ns
+            AND pw.checkpoint_id = checkpoints.checkpoint_id
+        ) as pending_writes,
+        (
+          SELECT
+            json_group_array(
+              json_object(
+                'type', ps.type,
+                'value', CAST(ps.value AS TEXT)
+              )
+            )
+          FROM writes as ps
+          WHERE ps.thread_id = checkpoints.thread_id
+            AND ps.checkpoint_ns = checkpoints.checkpoint_ns
+            AND ps.checkpoint_id = checkpoints.parent_checkpoint_id
+            AND ps.channel = '${TASKS}'
+          ORDER BY ps.idx
+        ) as pending_sends
+      FROM checkpoints\n`;
+
+    const whereClause: string[] = [];
+
+    if (thread_id) {
+      whereClause.push('thread_id = ?');
+    }
+
+    if (checkpoint_ns !== undefined && checkpoint_ns !== null) {
+      whereClause.push('checkpoint_ns = ?');
+    }
+
+    if (before?.configurable?.checkpoint_id !== undefined) {
+      whereClause.push('checkpoint_id < ?');
+    }
+
+    const sanitizedFilter = Object.fromEntries(
+      Object.entries(filter ?? {}).filter(
+        ([key, value]) =>
+          value !== undefined &&
+          validCheckpointMetadataKeys.includes(key as keyof CheckpointMetadata),
+      ),
+    );
+
+    whereClause.push(
+      ...Object.entries(sanitizedFilter).map(
+        ([key]) => `jsonb(CAST(metadata AS TEXT))->'$.${key}' = ?`,
+      ),
+    );
+
+    if (whereClause.length > 0) {
+      sql += `WHERE\n  ${whereClause.join(' AND\n  ')}\n`;
+    }
+
+    sql += '\nORDER BY checkpoint_id DESC';
+
+    if (limit) {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      sql += ` LIMIT ${parseInt(limit as any, 10)}`; // parseInt here (with cast to make TS happy) to sanitize input, as limit may be user-provided
+    }
+
+    const args = [
+      thread_id,
+      checkpoint_ns,
+      before?.configurable?.checkpoint_id,
+      ...Object.values(sanitizedFilter).map((value) => JSON.stringify(value)),
+    ].filter((value) => value !== undefined && value !== null);
+
+    const rows: CheckpointRow[] = this.db
+      .prepare(sql)
+      .all(...args) as CheckpointRow[];
+
+    if (rows) {
+      for (const row of rows) {
+        const pendingWrites = await Promise.all(
+          (JSON.parse(row.pending_writes) as PendingWriteColumn[]).map(
+            async (write) => {
+              return [
+                write.task_id,
+                write.channel,
+                await this.serde.loadsTyped(
+                  write.type ?? 'json',
+                  write.value ?? '',
+                ),
+              ] as [string, string, unknown];
+            },
+          ),
+        );
+
+        const messages = this.getMessageStmt.all(
+          row.thread_id,
+          row.checkpoint_ns,
+          row.checkpoint_id,
+        ) as MessageRow[];
+        const parsedMessages: BaseMessage[] = await Promise.all(
+          messages.map(async (message) => {
+            return this.serde.loadsTyped('json', message.message);
+          }),
+        );
+
+        const checkpoint = (await this.serde.loadsTyped(
+          row.type ?? 'json',
+          row.checkpoint,
+        )) as Checkpoint;
+
+        if (parsedMessages.length > 0) {
+          checkpoint.channel_values.messages = parsedMessages;
+        }
+
+        if (checkpoint.v < 4 && row.parent_checkpoint_id != null) {
+          await this.migratePendingSends(
+            checkpoint,
+            row.thread_id,
+            row.parent_checkpoint_id,
+          );
+        }
+
+        yield {
+          config: {
+            configurable: {
+              thread_id: row.thread_id,
+              checkpoint_ns: row.checkpoint_ns,
+              checkpoint_id: row.checkpoint_id,
+            },
+          },
+          checkpoint,
+          metadata: (await this.serde.loadsTyped(
+            row.type ?? 'json',
+            row.metadata,
+          )) as CheckpointMetadata,
+          parentConfig: row.parent_checkpoint_id
+            ? {
+                configurable: {
+                  thread_id: row.thread_id,
+                  checkpoint_ns: row.checkpoint_ns,
+                  checkpoint_id: row.parent_checkpoint_id,
+                },
+              }
+            : undefined,
+          pendingWrites,
+        };
+      }
+    }
+  }
+
+  async put(
+    config: RunnableConfig,
+    _checkpoint: Checkpoint,
+    metadata: CheckpointMetadata,
+  ): Promise<RunnableConfig> {
+    this.setup();
+
+    if (!config.configurable) {
+      throw new Error('Empty configuration supplied.');
+    }
+
+    const thread_id = config.configurable?.thread_id;
+    const checkpoint_ns = config.configurable?.checkpoint_ns ?? '';
+    const parent_checkpoint_id = config.configurable?.checkpoint_id;
+
+    if (!thread_id) {
+      throw new Error(
+        `Missing "thread_id" field in passed "config.configurable".`,
+      );
+    }
+
+    const { checkpoint, messages } = removeMessagesFromCheckpoint(_checkpoint);
+
+    const [[type1, serializedCheckpoint], [type2, serializedMetadata]] =
+      await Promise.all([
+        this.serde.dumpsTyped(checkpoint),
+        this.serde.dumpsTyped(metadata),
+      ]);
+
+    if (type1 !== type2) {
+      throw new Error(
+        'Failed to serialized checkpoint and metadata to the same type.',
+      );
+    }
+    const row = [
+      thread_id,
+      checkpoint_ns,
+      checkpoint.id,
+      parent_checkpoint_id,
+      type1,
+      serializedCheckpoint,
+      serializedMetadata,
+    ];
+
+    const transaction = this.db.transaction(() => {
+      this.putCheckpointStmt.run(...row);
+      if (messages) {
+        for (const message of messages) {
+          const encoder = new TextEncoder();
+          const msgFromMatrixRoom = message.additional_kwargs
+            .msgFromMatrixRoom as boolean;
+          const _additionalKwargs = message.additional_kwargs;
+          const cleanedAdditionalKwargs = cleanAdditionalKwargs(
+            message.additional_kwargs,
+            msgFromMatrixRoom ?? false,
+          );
+
+          message.additional_kwargs = {
+            ...cleanedAdditionalKwargs,
+            reasoning:
+              cleanedAdditionalKwargs.reasoning ?? _additionalKwargs.reasoning,
+            reasoningDetails:
+              cleanedAdditionalKwargs.reasoningDetails ??
+              _additionalKwargs.reasoningDetails,
+          };
+
+          if (message.type !== 'ai') {
+            delete message.additional_kwargs.reasoning;
+            delete message.additional_kwargs.reasoningDetails;
+          }
+
+          const serializedMessage = encoder.encode(
+            stringify(message, (_: string, value: any) => {
+              return _default(value);
+            }),
+          );
+
+          const messageRow: MessageRow = {
+            thread_id,
+            checkpoint_ns,
+            checkpoint_id: checkpoint.id,
+            message_id: message.id ?? message.lc_kwargs?.id,
+            message_type: message.type,
+            message_content: message.content.toString(),
+            message: serializedMessage as unknown as string,
+            created_at:
+              (message.additional_kwargs as CleanAdditionalKwargs)?.timestamp ??
+              new Date().toISOString(),
+          };
+
+          this.putMessageStmt.run(
+            messageRow.thread_id,
+            messageRow.checkpoint_ns,
+            messageRow.checkpoint_id,
+            messageRow.message_id,
+            messageRow.message_type,
+            messageRow.message_content,
+            messageRow.message,
+            messageRow.created_at,
+          );
+        }
+      }
+    });
+    transaction();
+
+    return {
+      configurable: {
+        thread_id,
+        checkpoint_ns,
+        checkpoint_id: checkpoint.id,
+      },
+    };
+  }
+
+  async putWrites(
+    config: RunnableConfig,
+    writes: PendingWrite[],
+    taskId: string,
+  ): Promise<void> {
+    this.setup();
+
+    if (!config.configurable) {
+      throw new Error('Empty configuration supplied.');
+    }
+
+    if (!config.configurable?.thread_id) {
+      console.error('Missing thread_id field in config.configurable.', {
+        configurable: config.configurable,
+      });
+
+      // get thread id using the checkpoint_id
+      const threadId = this.db
+        .prepare('SELECT thread_id FROM checkpoints WHERE checkpoint_id = ?')
+        .get(config.configurable?.checkpoint_id) as
+        | { thread_id: string }
+        | undefined;
+      if (!threadId) {
+        throw new Error(
+          'Missing thread_id field in config.configurable. config: ' +
+            JSON.stringify(config.configurable),
+        );
+      }
+      config.configurable.thread_id = threadId.thread_id;
+    }
+
+    if (!config.configurable?.checkpoint_id) {
+      console.error('Missing checkpoint_id field in config.configurable.', {
+        configurable: config.configurable,
+      });
+      throw new Error(
+        'Missing checkpoint_id field in config.configurable. config: ' +
+          JSON.stringify(config.configurable),
+      );
+    }
+
+    const transaction = this.db.transaction((rows) => {
+      for (const row of rows) {
+        this.putWritesStmt.run(...row);
+      }
+    });
+
+    const rows = await Promise.all(
+      writes.map(async (write, idx) => {
+        const [type, serializedWrite] = await this.serde.dumpsTyped(write[1]);
+        return [
+          config.configurable?.thread_id,
+          config.configurable?.checkpoint_ns,
+          config.configurable?.checkpoint_id,
+          taskId,
+          idx,
+          write[0],
+          type,
+          serializedWrite,
+        ];
+      }),
+    );
+
+    transaction(rows);
+  }
+
+  async deleteThread(threadId: string) {
+    const transaction = this.db.transaction(() => {
+      this.deleteCheckpointsStmt.run(threadId);
+      this.deleteWritesStmt.run(threadId);
+    });
+
+    transaction();
+  }
+
+  protected async migratePendingSends(
+    checkpoint: Checkpoint,
+    threadId: string,
+    parentCheckpointId: string,
+  ) {
+    const { pending_sends } = this.db
+      .prepare(
+        `
+          SELECT
+            checkpoint_id,
+            json_group_array(
+              json_object(
+                'type', ps.type,
+                'value', CAST(ps.value AS TEXT)
+              )
+            ) as pending_sends
+          FROM writes as ps
+          WHERE ps.thread_id = ?
+            AND ps.checkpoint_id = ?
+            AND ps.channel = '${TASKS}'
+          ORDER BY ps.idx
+        `,
+      )
+      .get(threadId, parentCheckpointId) as { pending_sends: string };
+
+    const mutableCheckpoint = checkpoint;
+
+    // add pending sends to checkpoint
+    mutableCheckpoint.channel_values ??= {};
+    mutableCheckpoint.channel_values[TASKS] = await Promise.all(
+      JSON.parse(pending_sends).map(({ type, value }: PendingSendColumn) =>
+        this.serde.loadsTyped(type, value),
+      ),
+    );
+
+    // add to versions
+    mutableCheckpoint.channel_versions[TASKS] =
+      Object.keys(checkpoint.channel_versions).length > 0
+        ? maxChannelVersion(...Object.values(checkpoint.channel_versions))
+        : this.getNextVersion(undefined);
+  }
+}
+
+const isCheckpointWithMessages = (
+  checkpoint: Checkpoint,
+): checkpoint is CheckpointWithMessages => {
+  return 'messages' in checkpoint.channel_values;
+};
+
+const removeMessagesFromCheckpoint = (
+  checkpoint: Checkpoint,
+): {
+  checkpoint: Checkpoint;
+  messages?: BaseMessage[];
+} => {
+  if (isCheckpointWithMessages(checkpoint)) {
+    const newCheckpoint = copyCheckpoint(checkpoint);
+    delete newCheckpoint.channel_values.messages;
+    return {
+      checkpoint: newCheckpoint,
+      messages: checkpoint.channel_values.messages,
+    };
+  }
+  return {
+    checkpoint,
+    messages: undefined,
+  };
+};