From c8ae5ddeed720cce07a9f31a766b2b215b1b82aa Mon Sep 17 00:00:00 2001
From: Ben Burns <803016+benjamincburns@users.noreply.github.com>
Date: Thu, 10 Oct 2024 19:14:16 +1300
Subject: [PATCH] fix: PostgresSaver doesn't store channel_values correctly

---
 libs/checkpoint-postgres/src/index.ts                    | 9 +++++----
 .../src/tests/checkpoints.int.test.ts                    | 6 ++----
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/libs/checkpoint-postgres/src/index.ts b/libs/checkpoint-postgres/src/index.ts
index 33818357..9b35e208 100644
--- a/libs/checkpoint-postgres/src/index.ts
+++ b/libs/checkpoint-postgres/src/index.ts
@@ -298,7 +298,9 @@ export class PostgresSaver extends BaseCheckpointSaver {
    * @param config The config to use for retrieving the checkpoint.
    * @returns The retrieved checkpoint tuple, or undefined.
    */
-  async getTuple(config: RunnableConfig): Promise<CheckpointTuple | undefined> {
+  async getTuple(
+    config: RunnableConfig
+  ): Promise<CheckpointTuple | undefined> {
     const {
       thread_id,
       checkpoint_ns = "",
@@ -416,8 +418,7 @@ export class PostgresSaver extends BaseCheckpointSaver {
   async put(
     config: RunnableConfig,
     checkpoint: Checkpoint,
-    metadata: CheckpointMetadata,
-    newVersions: ChannelVersions
+    metadata: CheckpointMetadata
   ): Promise<RunnableConfig> {
     if (config.configurable === undefined) {
       throw new Error(`Missing "configurable" field in "config" param`);
@@ -443,7 +444,7 @@ export class PostgresSaver extends BaseCheckpointSaver {
         thread_id,
         checkpoint_ns,
         checkpoint.channel_values,
-        newVersions
+        checkpoint.channel_versions
       );
       for (const serializedBlob of serializedBlobs) {
         await client.query(UPSERT_CHECKPOINT_BLOBS_SQL, serializedBlob);
diff --git a/libs/checkpoint-postgres/src/tests/checkpoints.int.test.ts b/libs/checkpoint-postgres/src/tests/checkpoints.int.test.ts
index fa08d603..3561f2e5 100644
--- a/libs/checkpoint-postgres/src/tests/checkpoints.int.test.ts
+++ b/libs/checkpoint-postgres/src/tests/checkpoints.int.test.ts
@@ -112,8 +112,7 @@ describe("PostgresSaver", () => {
     const runnableConfig = await postgresSaver.put(
       { configurable: { thread_id: "1" } },
       checkpoint1,
-      { source: "update", step: -1, writes: null, parents: {} },
-      checkpoint1.channel_versions
+      { source: "update", step: -1, writes: null, parents: {} }
     );
     expect(runnableConfig).toEqual({
       configurable: {
@@ -168,8 +167,7 @@ describe("PostgresSaver", () => {
         },
       },
       checkpoint2,
-      { source: "update", step: -1, writes: null, parents: {} },
-      checkpoint2.channel_versions
+      { source: "update", step: -1, writes: null, parents: {} }
     );
 
     // verify that parentTs is set and retrieved correctly for second checkpoint